Add apmath algorithms and list expression kind (#112)

Author: pearu

functional_algorithms/algorithms.py

@@ -13,6 +13,7 @@
 
 
 import functools
+import warnings
 from . import floating_point_algorithms as fpa
 
 
@@ -45,6 +46,16 @@ def foo(ctx, ...):
           # domain=None)` implements dispatch to real_foo and
           # complex_foo based on the arguments domain.
           assert 0  # unreachable
+
+    Warning: definition registry is global. When reusing definition
+    class in other modules, make sure to use module specific registry
+    via::
+
+      import functional_algorithms as fa
+
+      class definition(fa.algorithms.definition):
+        _registry = {}
+
     """
 
     # dict(<domain>=<dict of <native function name>:<definition for domain>>)
@@ -82,8 +93,13 @@ def wrapper(ctx, *args, **kwargs):
             if result is NotImplemented:
                 raise NotImplementedError(f"{self.native_func_name} not implemented for {self.domain} domain: {func.__name__}")
 
+            if isinstance(result, list):
+                result = ctx.list(result)
+
             return result
 
+        if self.native_func_name in self.registry:
+            warnings.warn(f"{self.native_func_name} wrapper is overwritten by {func.__name__} wrapper")
         self.registry[self.native_func_name] = wrapper
 
         return wrapper

functional_algorithms/apmath.py

@@ -187,7 +187,7 @@ def renormalize(ctx, seq, functional=False, fast=False, size=None, dtype=None):
     e_lst = vecsum(ctx, seq, fast=fast)
     # VecSumErrBranch:
     if functional:
-        zero = ctx.constant(0)
+        zero = ctx.constant(0, seq[0])
 
     f_lst = []
     eps_i = e_lst[0]
@@ -241,20 +241,21 @@ def nztopk(ctx, seq, k):
     elif len(seq) == 1:
         return seq
     elif len(seq) == 2:
-        flag = seq[0] == ctx.constant(0)
+        flag = seq[0] == ctx.constant(0, seq[0])
         if k == 1:
             return [ctx.select(flag, seq[1], seq[0])]
         return [ctx.select(flag, seq[1], seq[0]), ctx.select(flag, seq[0], seq[1])]
 
-    result = []
-    zero = ctx.constant(0)
-    one = ctx.constant(1)
+    zero = ctx.constant(0, seq[0])
+    izero = ctx.constant(0)
+    ione = ctx.constant(1, izero)
 
     isnzero = [a != zero for a in seq]
-    nzcount = [zero]  # ideally, nzcount ought to be a integer sequence
+    nzcount = [izero]
     for b in isnzero[:-1]:
-        nzcount.append(nzcount[-1] + ctx.select(b, one, zero))
+        nzcount.append(nzcount[-1] + ctx.select(b, ione, izero))
 
+    result = []
     for i in range(min(k, len(seq))):
         lst = [ctx.select(ctx.logical_and(isnzero[j], nzcount[j] == i), seq[j], zero) for j in range(i, len(seq))]
         result.append(sum(lst[:-1], lst[-1]))

functional_algorithms/apmath_algorithms.py

@@ -0,0 +1,18 @@
+from . import algorithms as faa
+from . import apmath
+
+
+class definition(faa.definition):
+    _registry = {}
+
+
+@definition("square", domain="real")
+def real_square(ctx, x: list[float, ...], functional: bool = True, size: int = None):
+    """Square on real input: x * x"""
+    return apmath.square(ctx, x, functional=functional, size=2)
+
+
+@definition("square")
+def square(ctx, z: list[float | complex, ...]):
+    """Square on floating-point expansion"""
+    assert 0

functional_algorithms/context.py

@@ -4,8 +4,8 @@
 import typing
 import warnings
 from collections import defaultdict
-from .utils import UNSPECIFIED, boolean_types, float_types, complex_types
-from .expr import Expr, make_constant, make_symbol, make_apply, known_expression_kinds
+from .utils import UNSPECIFIED, boolean_types, float_types, complex_types, integer_types
+from .expr import Expr, make_constant, make_symbol, make_apply, known_expression_kinds, make_list, make_item, make_len
 from .typesystem import Type
 
 
@@ -194,11 +194,14 @@ def trace(self, func, *args):
                     if ":" in a:
                         name, annot = a.split(":", 1)
                         name = name.strip()
+                        # TODO: parse annot as it may contain typing alias
                         param = param.replace(name=name if name else param.name, annotation=annot.strip())
                     else:
                         param = param.replace(name=a.strip())
                 elif isinstance(a, type):
                     param = param.replace(annotation=a.__name__)
+                elif isinstance(a, types.GenericAlias):
+                    param = param.replace(annotation=a)
                 else:
                     raise NotImplementedError((a, type(a)))
             if param.annotation is inspect.Parameter.empty:
@@ -207,12 +210,28 @@ def trace(self, func, *args):
                 typ = param.annotation
                 if isinstance(typ, types.UnionType):
                     typ = typing.get_args(typ)[0]
-                assert isinstance(typ, (type, str)), typ
+                assert isinstance(typ, (type, str, types.GenericAlias)), (type(typ), typ)
             if default_typ is UNSPECIFIED:
                 default_typ = typ
-            a = self.symbol(param.name, typ).reference(ref_name=param.name)
+            if isinstance(typ, types.GenericAlias):
+                if typ.__name__ == "list":
+                    a = self.list(
+                        [
+                            self.symbol(f"{param.name}_{k}_", t).reference(
+                                ref_name=f"{param.name}_{k}_",
+                                force=False,  # reference to item will be defined only when it is used
+                            )
+                            for k, t in enumerate(typ.__args__)
+                        ]
+                    )
+                else:
+                    raise TypeError(f"annotation type must be type of a scalar or list, got {typ}")
+            else:
+                a = self.symbol(param.name, typ)
+            a = a.reference(ref_name=param.name)
             new_args.append(a)
         args = tuple(new_args)
+
         name = self.symbol(func.__name__).reference(ref_name=func.__name__)
         return make_apply(self, name, args, func(self, *args))
 
@@ -233,6 +252,7 @@ def symbol(self, name, typ=UNSPECIFIED):
         if typ is UNSPECIFIED:
             like = self.default_like
             if like is not None:
+                assert like.kind == "symbol", like.kind
                 typ = like.operands[1]
             else:
                 typ = "float"
@@ -243,15 +263,31 @@ def constant(self, value, like_expr=UNSPECIFIED):
             if isinstance(value, boolean_types):
                 like_expr = self.symbol("_boolean_value", "boolean")
             elif self._default_constant_type is not None:
+                # Warning: when specified and default_constant_type is
+                # float, integer values will be interpreted as floats
                 like_expr = self.symbol("_value", self._default_constant_type)
+            elif isinstance(value, integer_types):
+                like_expr = self.symbol("_integer_value", type(value))
             elif isinstance(value, float_types):
                 like_expr = self.symbol("_float_value", type(value))
             elif isinstance(value, complex_types):
                 like_expr = self.symbol("_complex_value", type(value))
             else:
                 like_expr = self.default_like
+        elif isinstance(like_expr, (str, type, Type)):
+            typ = Type.fromobject(self, like_expr)
+            like_expr = self.symbol(f"_{typ.kind}_value", typ)
         return make_constant(self, value, like_expr)
 
+    def list(self, items):
+        return make_list(self, items)
+
+    def item(self, container, index):
+        return make_item(self, container, index)
+
+    def len(self, container):
+        return make_len(self, container)
+
     def call(self, func, args):
         """Apply callable to arguments and return its result.