MAINT: rework unary and binary ufuncs w/ normalizations

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -1,52 +1,60 @@
-from ._decorators import deco_binary_ufunc_from_impl
-from ._detail import _ufunc_impl
+from ._detail import _binary_ufuncs
+
+__all__ = [name for name in dir(_binary_ufuncs) if not name.startswith("_") and name != "torch"]
+
+
+# TODO: consolidate normalizations
+from ._funcs import normalizer, ArrayLike, SubokLike, DTypeLike
+from ._detail import _util
+from . import _helpers
+
+
+def deco_binary_ufunc(torch_func):
+    """Common infra for unary ufuncs.
+
+    Normalize arguments, sort out type casting, broadcasting and delegate to
+    the pytorch functions for the actual work.
+    """
+    def wrapped(
+        x1 : ArrayLike,
+        x2 : ArrayLike,
+        /,
+        out=None,
+        *,
+        where=True,
+        casting="same_kind",
+        order="K",
+        dtype: DTypeLike=None,
+        subok: SubokLike=False,
+        signature=None,
+        extobj=None
+    ):
+        if order != "K" or not where or signature or extobj:
+            raise NotImplementedError
+
+        # XXX: dtype=... parameter
+        if dtype is not None:
+            raise NotImplementedError
+
+        out_shape_dtype = None
+        if out is not None:
+            out_shape_dtype = (out.get().dtype, out.get().shape)
+
+        tensors = _util.cast_and_broadcast((x1, x2), out_shape_dtype, casting)
+
+        result = torch_func(*tensors)
+        return _helpers.result_or_out(result, out)
+
+    return wrapped
 
 #
-# Functions in this file implement binary ufuncs: wrap two first arguments in
-# asarray and delegate to functions from _ufunc_impl.
-#
-# Functions in _detail/_ufunc_impl.py receive tensors, implement common tasks
-# with ufunc args, and delegate heavy lifting to pytorch equivalents.
+# For each torch ufunc implementation, decorate and attach the decorated name
+# to this module. Its contents is then exported to the public namespace in __init__.py
 #
+for name in __all__:
+    ufunc = getattr(_binary_ufuncs, name)
+    decorated = normalizer(deco_binary_ufunc(ufunc))
 
-# the list is autogenerated, cf autogen/gen_ufunc_2.py
-add = deco_binary_ufunc_from_impl(_ufunc_impl.add)
-arctan2 = deco_binary_ufunc_from_impl(_ufunc_impl.arctan2)
-bitwise_and = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_and)
-bitwise_or = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_or)
-bitwise_xor = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_xor)
-copysign = deco_binary_ufunc_from_impl(_ufunc_impl.copysign)
-divide = deco_binary_ufunc_from_impl(_ufunc_impl.divide)
-equal = deco_binary_ufunc_from_impl(_ufunc_impl.equal)
-float_power = deco_binary_ufunc_from_impl(_ufunc_impl.float_power)
-floor_divide = deco_binary_ufunc_from_impl(_ufunc_impl.floor_divide)
-fmax = deco_binary_ufunc_from_impl(_ufunc_impl.fmax)
-fmin = deco_binary_ufunc_from_impl(_ufunc_impl.fmin)
-fmod = deco_binary_ufunc_from_impl(_ufunc_impl.fmod)
-gcd = deco_binary_ufunc_from_impl(_ufunc_impl.gcd)
-greater = deco_binary_ufunc_from_impl(_ufunc_impl.greater)
-greater_equal = deco_binary_ufunc_from_impl(_ufunc_impl.greater_equal)
-heaviside = deco_binary_ufunc_from_impl(_ufunc_impl.heaviside)
-hypot = deco_binary_ufunc_from_impl(_ufunc_impl.hypot)
-lcm = deco_binary_ufunc_from_impl(_ufunc_impl.lcm)
-ldexp = deco_binary_ufunc_from_impl(_ufunc_impl.ldexp)
-left_shift = deco_binary_ufunc_from_impl(_ufunc_impl.left_shift)
-less = deco_binary_ufunc_from_impl(_ufunc_impl.less)
-less_equal = deco_binary_ufunc_from_impl(_ufunc_impl.less_equal)
-logaddexp = deco_binary_ufunc_from_impl(_ufunc_impl.logaddexp)
-logaddexp2 = deco_binary_ufunc_from_impl(_ufunc_impl.logaddexp2)
-logical_and = deco_binary_ufunc_from_impl(_ufunc_impl.logical_and)
-logical_or = deco_binary_ufunc_from_impl(_ufunc_impl.logical_or)
-logical_xor = deco_binary_ufunc_from_impl(_ufunc_impl.logical_xor)
-matmul = deco_binary_ufunc_from_impl(_ufunc_impl.matmul)
-maximum = deco_binary_ufunc_from_impl(_ufunc_impl.maximum)
-minimum = deco_binary_ufunc_from_impl(_ufunc_impl.minimum)
-remainder = deco_binary_ufunc_from_impl(_ufunc_impl.remainder)
-multiply = deco_binary_ufunc_from_impl(_ufunc_impl.multiply)
-nextafter = deco_binary_ufunc_from_impl(_ufunc_impl.nextafter)
-not_equal = deco_binary_ufunc_from_impl(_ufunc_impl.not_equal)
-power = deco_binary_ufunc_from_impl(_ufunc_impl.power)
-remainder = deco_binary_ufunc_from_impl(_ufunc_impl.remainder)
-right_shift = deco_binary_ufunc_from_impl(_ufunc_impl.right_shift)
-subtract = deco_binary_ufunc_from_impl(_ufunc_impl.subtract)
-divide = deco_binary_ufunc_from_impl(_ufunc_impl.divide)
+    decorated.__qualname__ = name    # XXX: is this really correct?
+    decorated.__name__ = name
+    vars()[name] = decorated

torch_np/_decorators.py

@@ -9,31 +9,6 @@
 NoValue = None
 
 
-def dtype_to_torch(func):
-    @functools.wraps(func)
-    def wrapped(*args, dtype=None, **kwds):
-        torch_dtype = None
-        if dtype is not None:
-            dtype = _dtypes.dtype(dtype)
-            torch_dtype = dtype.torch_dtype
-        return func(*args, dtype=torch_dtype, **kwds)
-
-    return wrapped
-
-
-def emulate_out_arg(func):
-    """Simulate the out=... handling: move the result tensor to the out array.
-
-    With this decorator, the inner function just does not see the out array.
-    """
-
-    @functools.wraps(func)
-    def wrapped(*args, out=None, **kwds):
-        result_tensor = func(*args, **kwds)
-        return _helpers.result_or_out(result_tensor, out)
-
-    return wrapped
-
 
 def out_shape_dtype(func):
     """Handle out=... kwarg for ufuncs.
@@ -53,32 +28,3 @@ def wrapped(*args, out=None, **kwds):
     return wrapped
 
 
-def deco_unary_ufunc_from_impl(impl_func):
-    @functools.wraps(impl_func)
-    @dtype_to_torch
-    @out_shape_dtype
-    def wrapped(x1, *args, **kwds):
-        from ._ndarray import asarray
-
-        x1_tensor = asarray(x1).get()
-        result = impl_func((x1_tensor,), *args, **kwds)
-        return result
-
-    return wrapped
-
-
-# TODO: deduplicate with _ndarray/asarray_replacer,
-# and _wrapper/concatenate et al
-def deco_binary_ufunc_from_impl(impl_func):
-    @functools.wraps(impl_func)
-    @dtype_to_torch
-    @out_shape_dtype
-    def wrapped(x1, x2, *args, **kwds):
-        from ._ndarray import asarray
-
-        x1_tensor = asarray(x1).get()
-        x2_tensor = asarray(x2).get()
-        return impl_func((x1_tensor, x2_tensor), *args, **kwds)
-
-    return wrapped
-

torch_np/_detail/_binary_ufuncs.py

@@ -0,0 +1,30 @@
+"""Export torch work functions for binary ufuncs, rename/tweak to match numpy.
+This listing is further exported to public symbols in the `torch_np/_binary_ufuncs.py` module.
+"""
+
+import torch
+
+from . import _dtypes_impl, _util
+
+
+from torch import (add, arctan2, bitwise_and, bitwise_or, bitwise_xor, copysign, divide, 
+    float_power, floor_divide, fmax, fmin, fmod, gcd, greater, greater_equal, heaviside,
+    hypot, lcm, ldexp, less, less_equal, logaddexp, logaddexp2, logical_and,
+    logical_or, logical_xor, maximum, minimum, remainder, multiply, nextafter, not_equal,
+    remainder, subtract, divide)
+
+
+# renames
+from torch import (eq as equal, pow as power, bitwise_right_shift as right_shift,
+        bitwise_left_shift as left_shift,)
+
+
+# work around torch limitations w.r.t. numpy
+def matmul(x, y):
+    # work around RuntimeError: expected scalar type Int but found Double
+    dtype = _dtypes_impl.result_type_impl((x.dtype, y.dtype))
+    x = _util.cast_if_needed(x, dtype)
+    y = _util.cast_if_needed(y, dtype)
+    result = torch.matmul(x, y)
+    return result
+

torch_np/_detail/_ufunc_impl.py

@@ -0,0 +1,35 @@
+"""Export torch work functions for unary ufuncs, rename/tweak to match numpy.
+This listing is further exported to public symbols in the `torch_np/_unary_ufuncs.py` module.
+"""
+
+import torch
+
+from torch import (arccos, arccosh, arcsin, arcsinh, arctan, arctanh, ceil,
+    cos, cosh, deg2rad, exp, exp2, expm1,
+    floor, isfinite, isinf, isnan, log, log10, log1p, log2, logical_not,
+    negative, rad2deg, reciprocal, sign, signbit,
+    sin, sinh, sqrt, square, tan, tanh, trunc)
+
+# renames
+from torch import (conj_physical as conjugate, round as rint, bitwise_not as invert, rad2deg as degrees,
+    deg2rad as radians, absolute as fabs, )
+
+# special cases: torch does not export these names
+def cbrt(x):
+    return torch.pow(x, 1 / 3)
+
+
+def positive(x):
+    return +x
+
+
+def absolute(x):
+    # work around torch.absolute not impl for bools
+    if x.dtype == torch.bool:
+        return x
+    return torch.absolute(x)
+
+
+abs = absolute
+conj = conjugate
+