Trigonometric argument reduction for float16

functional_algorithms/__init__.py

@@ -20,3 +20,4 @@
 from .restrict import restrict
 from .signatures import filter_signatures
 from . import fpu
+from .typesystem import Type

functional_algorithms/context.py

@@ -1,14 +1,37 @@
+import contextlib
 import inspect
 import sys
 import types
 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_series
 from .typesystem import Type
 
 
+class Parameters(dict):
+    """A dictionary that call method enables context manager support to
+    temporarily modify dictionary content.
+    """
+
+    def __call__(self, **items):
+
+        @contextlib.contextmanager
+        def manager(items):
+            prev_items = dict((k, v) for k, v in self.items() if k in items)
+            new_keys = {k for k in items if k not in self}
+            try:
+                self.update(items)
+                yield self
+            finally:
+                self.update(prev_items)
+                for k in new_keys:
+                    self.pop(k)
+
+        return manager(items)
+
+
 class Context:
 
     def __init__(self, paths=[], enable_alt=None, default_constant_type=None, parameters=None):
@@ -39,7 +62,7 @@ def __init__(self, paths=[], enable_alt=None, default_constant_type=None, parame
         self._enable_alt = enable_alt
         self._default_constant_type = default_constant_type
         self._default_like = None
-        self.parameters = parameters or {}
+        self.parameters = Parameters(parameters or {})
         if "using" not in self.parameters:
             self.parameters["using"] = set()
 
@@ -248,8 +271,12 @@ def constant(self, value, like_expr=UNSPECIFIED):
                 like_expr = self.symbol("_float_value", type(value))
             elif isinstance(value, complex_types):
                 like_expr = self.symbol("_complex_value", type(value))
+            elif isinstance(value, integer_types):
+                like_expr = self.symbol("_integer_value", type(value))
             else:
                 like_expr = self.default_like
+        if like_expr is int:
+            like_expr = self.symbol("_integer_value", int)
         return make_constant(self, value, like_expr)
 
     def call(self, func, args):
@@ -317,8 +344,10 @@ def floor_divide(self, x, y):
     def pow(self, x, y):
         if isinstance(y, float) and y == 0.5:
             return self.sqrt(x)
-        if isinstance(y, int) and y == 2:
-            return self.square(x)
+        if isinstance(y, int):
+            if y == 2:
+                return self.square(x)
+            return Expr(self, "pow", (x, self.constant(y, int)))
         return Expr(self, "pow", (x, y))
 
     def logical_and(self, x, y):
@@ -341,8 +370,8 @@ def logical_not(self, x):
 
     Not = logical_not
 
-    def bitwise_invert(self):
-        return Expr(self, "bitwise_invert", (self,))
+    def bitwise_invert(self, x):
+        return Expr(self, "bitwise_invert", (x,))
 
     invert = bitwise_invert
 
@@ -516,3 +545,19 @@ def downcast(self, x):
 
     def is_finite(self, x):
         return Expr(self, "is_finite", (x,))
+
+    def series(self, *terms, **params):
+        unit_index = params.get("unit_index", 0)
+        scaling_exp = params.get("scaling_exp", 0)
+        return make_series(self, unit_index, scaling_exp, terms)
+
+    def _series(self, terms, params):
+        return self.series(*terms, **params)
+
+    def _get_series_operands(self, expr):
+        if expr.kind == "series":
+            return expr.operands
+        return expr
+
+    def fma(self, x, y, z):
+        return Expr(self, "fma", (x, y, z))

functional_algorithms/expr.py

@@ -23,7 +23,8 @@
 ceil, floor, floor_divide, remainder, round, truncate,
 copysign, sign, nextafter,
 upcast, downcast,
-is_finite, is_inf, is_posinf, is_neginf, is_nan, is_negzero
+is_finite, is_inf, is_posinf, is_neginf, is_nan, is_negzero,
+series, fma
 """.replace(
         " ", ""
     )
@@ -80,6 +81,7 @@ def normalize_like(expr):
             "sqrt",
             "square",
             "asin_acos_kernel",
+            "fma",
         }:
             expr = expr.operands[0]
         elif expr.kind == "absolute" and not expr.operands[0].is_complex:
@@ -88,6 +90,8 @@ def normalize_like(expr):
             expr = expr.operands[0].operands[0]
         elif expr.kind == "imag" and expr.operands[0].kind == "complex":
             expr = expr.operands[0].operands[1]
+        elif expr.kind == "series":
+            expr = expr.operands[1]
         else:
             break
     return expr
@@ -118,20 +122,29 @@ def make_apply(context, name, args, result):
     return Expr(context, "apply", (name, *args, result))
 
 
+def make_series(context, unit_index, scaling_exp, terms):
+    return Expr(context, "series", ((unit_index, scaling_exp), *terms))
+
+
 def normalize(context, operands):
     """Convert numbers to constant expressions"""
     exprs = [operand for operand in operands if isinstance(operand, Expr)]
     if len(exprs) == 0:
         ref_operand = context.default_like
         if ref_operand is None:
-            raise ValueError("cannot normalize operands with no reference operand or context default constant type")
+            operand_types = ", ".join(sorted(set([type(operand).__name__ for operand in operands])))
+            raise ValueError(
+                f"cannot normalize operands (of types {operand_types}) with no reference operand or context default constant type"
+            )
     else:
         ref_operand = exprs[0]
 
     new_operands = []
     for operand in operands:
-        if isinstance(operand, (int, float, complex, str)):
+        if isinstance(operand, (int, float, complex, str, numpy.floating, numpy.complexfloating, numpy.integer)):
             operand = make_constant(context, operand, ref_operand)
+        else:
+            assert isinstance(operand, Expr), type(operand)
         new_operands.append(operand)
     return tuple(new_operands)
 
@@ -161,7 +174,7 @@ def toidentifier(value):
     elif isinstance(value, numpy.floating):
         try:
             intvalue = int(value)
-        except OverflowError:
+        except (OverflowError, ValueError):
             intvalue = None
         if value == intvalue:
             return value.dtype.kind + toidentifier(intvalue)
@@ -201,10 +214,19 @@ def make_ref(expr):
         elif expr.kind == "absolute":
             # using abs for BC
             ref = f"abs_{make_ref(expr.operands[0])}"
-        else:
+        elif expr.kind == "series":
             all_operands_have_ref_name = not [
-                0 for o in expr.operands if not isinstance(expr.operands[0].props.get("reference_name"), str)
+                0 for o in expr.operands[1:] if not isinstance(o.props.get("reference_name"), str)
             ]
+            if all_operands_have_ref_name:
+                # for readability
+                params = [f"minus{p}" if p < 0 else str(p) for p in expr.operands[0]]
+                lst = [expr.kind] + params + list(map(make_ref, expr.operands[1:]))
+                ref = "_".join(lst)
+            else:
+                ref = f"{expr.kind}_{expr.intkey}"
+        else:
+            all_operands_have_ref_name = not [0 for o in expr.operands if not isinstance(o.props.get("reference_name"), str)]
             if all_operands_have_ref_name:
                 # for readability
                 lst = [expr.kind] + list(map(make_ref, expr.operands))
@@ -315,7 +337,7 @@ def __new__(cls, context, kind, operands):
             if context.alt is not None and not isinstance(operands[0], Expr):
                 operands = [context.alt.constant(operands[0]), operands[1]]
         else:
-            if kind == "select":
+            if kind in {"select", "series"}:
                 operands = operands[:1] + normalize(context, operands[1:])
             else:
                 operands = normalize(context, operands)
@@ -338,7 +360,13 @@ def __new__(cls, context, kind, operands):
                     operands = (Expr(context.alt, kind, tuple(constant_operands)), constant_like)
                     kind = "constant"
 
-            assert False not in [isinstance(operand, Expr) for operand in operands], operands
+            if kind == "series":
+                assert isinstance(operands[0], tuple) and len(operands[0]) == 2, type(operands[0])
+                assert isinstance(operands[0][0], int)
+                assert isinstance(operands[0][1], int)
+                assert False not in [isinstance(operand, Expr) for operand in operands[1:]], operands
+            else:
+                assert False not in [isinstance(operand, Expr) for operand in operands], operands
 
         obj.context = context
         obj.kind = kind
@@ -387,6 +415,8 @@ def _compute_serialized(self):
                 value.key if isinstance(value, Expr) else (value, type(value).__name__),
                 like.key,
             )
+        elif self.kind == "series":
+            r = (self.kind, self.operands[0], *(operand.intkey for operand in self.operands[1:]))
         else:
             # Don't use `operand.key` as its size is unbounded and
             # will lead to large overhead in computing the hash value
@@ -467,6 +497,14 @@ def rewrite(self, modifier, *modifiers, deep_first=None, _rewrite_context=None):
                 result = make_apply(self.context, self.operands[0], self.operands[1:-1], body)
             else:
                 result = self
+        elif self.kind == "series":
+            terms = tuple([operand.rewrite(modifier, **rewrite_kwargs) for operand in self.operands[1:]])
+            for o1, o2 in zip(terms, self.operands[1:]):
+                if o1 is not o2:
+                    result = Expr(self.context, self.kind, (self.operands[0], *terms))
+                    break
+            else:
+                result = self
         else:
             operands = tuple([operand.rewrite(modifier, **rewrite_kwargs) for operand in self.operands])
             for o1, o2 in zip(operands, self.operands):
@@ -584,8 +622,13 @@ def __str__(self):
         return Printer().tostring(self)
 
     def __repr__(self):
-        operands = tuple(f"{o.kind}:{o.intkey}" for o in self.operands)
-        return f"{type(self).__name__}({self.kind}, {operands}, {self.props})"
+        if self.kind in {"symbol"}:
+            operands = tuple(repr(o) for o in self.operands)
+        elif self.kind in {"series", "constant"}:
+            operands = (repr(self.operands[0]),) + tuple(f"{o.kind}:{o.intkey}" for o in self.operands[1:])
+        else:
+            operands = tuple(f"{o.kind}:{o.intkey}" for o in self.operands)
+        return f"{type(self).__name__}({self.kind!r}, {operands}, {self.props})"
 
     def __abs__(self):
         return self.context.absolute(self)
@@ -989,12 +1032,16 @@ def is_complex(self):
             "minimum",
             "floor_divide",
             "remainder",
+            "truncate",
+            "round",
         }:
             return False
         elif self.kind in {"complex", "conjugate"}:
             return True
         elif self.kind in {"add", "subtract", "divide", "multiply", "pow"}:
             return self.operands[0].is_complex or self.operands[1].is_complex
+        elif self.kind == "fma":
+            return self.operands[0].is_complex or self.operands[1].is_complex or self.operands[2].is_complex
         elif self.kind in {
             "positive",
             "negative",
@@ -1023,6 +1070,11 @@ def is_complex(self):
             return self.operands[0].is_complex
         elif self.kind == "apply":
             return self.operands[-1].is_complex
+        elif self.kind == "series":
+            for o in self.operands[1:]:
+                if o.is_complex:
+                    return True
+            return False
         else:
             raise NotImplementedError(f"{type(self).__name__}.is_complex not implemented for {self.kind}")
 
@@ -1065,6 +1117,8 @@ def get_type(self):
             "copysign",
             "conjugate",
             "asin_acos_kernel",
+            "truncate",
+            "round",
         }:
             return self.operands[0].get_type()
         elif self.kind in {
@@ -1080,6 +1134,16 @@ def get_type(self):
             "atan2",
         }:
             return self.operands[0].get_type().max(self.operands[1].get_type())
+        elif self.kind == "series":
+            t = self.operands[1].get_type()
+            for o in self.operands[2:]:
+                t = t.max(o.get_type())
+            return t
+        elif self.kind == "fma":
+            t = self.operands[0].get_type()
+            for o in self.operands[1:]:
+                t = t.max(o.get_type())
+            return t
         elif self.kind in {"absolute", "real", "imag"}:
             t = self.operands[0].get_type()
             return t.complex_part if t.is_complex else t