Merge pull request #70 from Quansight-Labs/normalizations

bare-bones normalizations via type hints

Author: ev-br

torch_np/__init__.py

@@ -1,4 +1,3 @@
-from ._wrapper import *  # isort: skip  # XXX: currently this prevents circular imports
 from . import random
 from ._binary_ufuncs import *
 from ._detail._index_tricks import *
@@ -8,6 +7,7 @@
 from ._getlimits import errstate, finfo, iinfo
 from ._ndarray import array, asarray, can_cast, ndarray, newaxis, result_type
 from ._unary_ufuncs import *
+from ._wrapper import *
 
 # from . import testing
 

torch_np/_binary_ufuncs.py

@@ -1,52 +1,52 @@
-from ._decorators import deco_binary_ufunc_from_impl
-from ._detail import _ufunc_impl
+from typing import Optional
+
+from . import _helpers
+from ._detail import _binary_ufuncs
+from ._normalizations import ArrayLike, DTypeLike, NDArray, SubokLike, normalizer
+
+__all__ = [
+    name for name in dir(_binary_ufuncs) if not name.startswith("_") and name != "torch"
+]
+
+
+def deco_binary_ufunc(torch_func):
+    """Common infra for binary 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: Optional[NDArray] = None,
+        *,
+        where=True,
+        casting="same_kind",
+        order="K",
+        dtype: DTypeLike = None,
+        subok: SubokLike = False,
+        signature=None,
+        extobj=None,
+    ):
+        tensors = _helpers.ufunc_preprocess(
+            (x1, x2), out, where, casting, order, dtype, subok, signature, extobj
+        )
+        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

@@ -1,39 +1,10 @@
 import functools
-import operator
 
 import torch
 
 from . import _dtypes, _helpers
 from ._detail import _util
 
-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.
@@ -51,89 +22,3 @@ def wrapped(*args, out=None, **kwds):
         return _helpers.result_or_out(result_tensor, out)
 
     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
-
-
-def axis_keepdims_wrapper(func):
-    """`func` accepts an array-like as a 1st arg, returns a tensor.
-
-    This decorator implements the generic handling of axis, out and keepdims
-    arguments for reduction functions.
-
-    Note that we peel off `out=...` and `keepdims=...` args (torch functions never
-    see them). The `axis` argument we normalize and pass through to pytorch functions.
-
-    """
-    # TODO: sort out function signatures: how they flow through all decorators etc
-    @functools.wraps(func)
-    def wrapped(a, axis=None, keepdims=NoValue, *args, **kwds):
-        from ._ndarray import asarray, ndarray
-
-        tensor = asarray(a).get()
-
-        # standardize the axis argument
-        if isinstance(axis, ndarray):
-            axis = operator.index(axis)
-
-        result = _util.axis_expand_func(func, tensor, axis, *args, **kwds)
-
-        if keepdims:
-            result = _util.apply_keepdims(result, axis, tensor.ndim)
-
-        return result
-
-    return wrapped
-
-
-def axis_none_ravel_wrapper(func):
-    """`func` accepts an array-like as a 1st arg, returns a tensor.
-
-    This decorator implements the generic handling of axis=None acting on a
-    raveled array. One use is cumprod / cumsum. concatenate also uses a
-    similar logic.
-
-    """
-
-    @functools.wraps(func)
-    def wrapped(a, axis=None, *args, **kwds):
-        from ._ndarray import asarray, ndarray
-
-        tensor = asarray(a).get()
-
-        # standardize the axis argument
-        if isinstance(axis, ndarray):
-            axis = operator.index(axis)
-
-        result = _util.axis_ravel_func(func, tensor, axis, *args, **kwds)
-        return result
-
-    return wrapped

torch_np/_detail/_binary_ufuncs.py

@@ -0,0 +1,53 @@
+"""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
+
+# renames
+from torch import add, arctan2, bitwise_and
+from torch import bitwise_left_shift as left_shift
+from torch import bitwise_or
+from torch import bitwise_right_shift as right_shift
+from torch import bitwise_xor, copysign, divide
+from torch import eq as equal
+from torch import (
+    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,
+    multiply,
+    nextafter,
+    not_equal,
+)
+from torch import pow as power
+from torch import remainder, subtract
+
+from . import _dtypes_impl, _util
+
+
+# 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