MAINT: merge {_unary,_binary}_ufuncs modules

Author: ev-br

torch_np/__init__.py

@@ -1,11 +1,10 @@
 from . import linalg, random
-from ._binary_ufuncs import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
 from ._funcs import *
 from ._getlimits import errstate, finfo, iinfo
 from ._ndarray import array, asarray, can_cast, ndarray, newaxis, result_type
-from ._unary_ufuncs import *
+from ._ufuncs import *
 
 # from . import testing
 

torch_np/_detail/_binary_ufuncs.py renamed to torch_np/_binary_ufuncs_impl.py

@@ -1,5 +1,5 @@
 """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.
+This listing is further exported to public symbols in the `torch_np/_ufuncs.py` module.
 """
 
 import torch
@@ -42,7 +42,7 @@
 from torch import remainder as mod
 from torch import subtract, true_divide
 
-from . import _dtypes_impl, _util
+from ._detail import _dtypes_impl, _util
 
 
 # work around torch limitations w.r.t. numpy

torch_np/_ndarray.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from . import _binary_ufuncs, _dtypes, _funcs, _funcs_impl, _helpers, _unary_ufuncs
+from . import _dtypes, _funcs, _funcs_impl, _helpers, _ufuncs
 from ._detail import _dtypes_impl, _util
 
 newaxis = None
@@ -176,24 +176,24 @@ def __str__(self):
     ### comparisons ###
     def __eq__(self, other):
         try:
-            return _binary_ufuncs.equal(self, other)
+            return _ufuncs.equal(self, other)
         except (RuntimeError, TypeError):
             # Failed to convert other to array: definitely not equal.
             falsy = torch.full(self.shape, fill_value=False, dtype=bool)
             return asarray(falsy)
 
     def __ne__(self, other):
         try:
-            return _binary_ufuncs.not_equal(self, other)
+            return _ufuncs.not_equal(self, other)
         except (RuntimeError, TypeError):
             # Failed to convert other to array: definitely not equal.
             falsy = torch.full(self.shape, fill_value=True, dtype=bool)
             return asarray(falsy)
 
-    __gt__ = _binary_ufuncs.greater
-    __lt__ = _binary_ufuncs.less
-    __ge__ = _binary_ufuncs.greater_equal
-    __le__ = _binary_ufuncs.less_equal
+    __gt__ = _ufuncs.greater
+    __lt__ = _ufuncs.less
+    __ge__ = _ufuncs.greater_equal
+    __le__ = _ufuncs.less_equal
 
     def __bool__(self):
         try:
@@ -239,107 +239,107 @@ def __len__(self):
     ### arithmetic ###
 
     # add, self + other
-    __add__ = __radd__ = _binary_ufuncs.add
+    __add__ = __radd__ = _ufuncs.add
 
     def __iadd__(self, other):
-        return _binary_ufuncs.add(self, other, out=self)
+        return _ufuncs.add(self, other, out=self)
 
     # sub, self - other
-    __sub__ = _binary_ufuncs.subtract
+    __sub__ = _ufuncs.subtract
 
     # XXX: generate a function just for this? AND other non-commutative ops.
     def __rsub__(self, other):
-        return _binary_ufuncs.subtract(other, self)
+        return _ufuncs.subtract(other, self)
 
     def __isub__(self, other):
-        return _binary_ufuncs.subtract(self, other, out=self)
+        return _ufuncs.subtract(self, other, out=self)
 
     # mul, self * other
-    __mul__ = __rmul__ = _binary_ufuncs.multiply
+    __mul__ = __rmul__ = _ufuncs.multiply
 
     def __imul__(self, other):
-        return _binary_ufuncs.multiply(self, other, out=self)
+        return _ufuncs.multiply(self, other, out=self)
 
     # div, self / other
-    __truediv__ = _binary_ufuncs.divide
+    __truediv__ = _ufuncs.divide
 
     def __rtruediv__(self, other):
-        return _binary_ufuncs.divide(other, self)
+        return _ufuncs.divide(other, self)
 
     def __itruediv__(self, other):
-        return _binary_ufuncs.divide(self, other, out=self)
+        return _ufuncs.divide(self, other, out=self)
 
     # floordiv, self // other
-    __floordiv__ = _binary_ufuncs.floor_divide
+    __floordiv__ = _ufuncs.floor_divide
 
     def __rfloordiv__(self, other):
-        return _binary_ufuncs.floor_divide(other, self)
+        return _ufuncs.floor_divide(other, self)
 
     def __ifloordiv__(self, other):
-        return _binary_ufuncs.floor_divide(self, other, out=self)
+        return _ufuncs.floor_divide(self, other, out=self)
 
-    __divmod__ = _binary_ufuncs.divmod
+    __divmod__ = _ufuncs.divmod
 
     # power, self**exponent
-    __pow__ = __rpow__ = _binary_ufuncs.float_power
+    __pow__ = __rpow__ = _ufuncs.float_power
 
     def __rpow__(self, exponent):
-        return _binary_ufuncs.float_power(exponent, self)
+        return _ufuncs.float_power(exponent, self)
 
     def __ipow__(self, exponent):
-        return _binary_ufuncs.float_power(self, exponent, out=self)
+        return _ufuncs.float_power(self, exponent, out=self)
 
     # remainder, self % other
-    __mod__ = __rmod__ = _binary_ufuncs.remainder
+    __mod__ = __rmod__ = _ufuncs.remainder
 
     def __imod__(self, other):
-        return _binary_ufuncs.remainder(self, other, out=self)
+        return _ufuncs.remainder(self, other, out=self)
 
     # bitwise ops
     # and, self & other
-    __and__ = __rand__ = _binary_ufuncs.bitwise_and
+    __and__ = __rand__ = _ufuncs.bitwise_and
 
     def __iand__(self, other):
-        return _binary_ufuncs.bitwise_and(self, other, out=self)
+        return _ufuncs.bitwise_and(self, other, out=self)
 
     # or, self | other
-    __or__ = __ror__ = _binary_ufuncs.bitwise_or
+    __or__ = __ror__ = _ufuncs.bitwise_or
 
     def __ior__(self, other):
-        return _binary_ufuncs.bitwise_or(self, other, out=self)
+        return _ufuncs.bitwise_or(self, other, out=self)
 
     # xor, self ^ other
-    __xor__ = __rxor__ = _binary_ufuncs.bitwise_xor
+    __xor__ = __rxor__ = _ufuncs.bitwise_xor
 
     def __ixor__(self, other):
-        return _binary_ufuncs.bitwise_xor(self, other, out=self)
+        return _ufuncs.bitwise_xor(self, other, out=self)
 
     # bit shifts
-    __lshift__ = __rlshift__ = _binary_ufuncs.left_shift
+    __lshift__ = __rlshift__ = _ufuncs.left_shift
 
     def __ilshift__(self, other):
-        return _binary_ufuncs.left_shift(self, other, out=self)
+        return _ufuncs.left_shift(self, other, out=self)
 
-    __rshift__ = __rrshift__ = _binary_ufuncs.right_shift
+    __rshift__ = __rrshift__ = _ufuncs.right_shift
 
     def __irshift__(self, other):
-        return _binary_ufuncs.right_shift(self, other, out=self)
+        return _ufuncs.right_shift(self, other, out=self)
 
-    __matmul__ = _binary_ufuncs.matmul
+    __matmul__ = _ufuncs.matmul
 
     def __rmatmul__(self, other):
-        return _binary_ufuncs.matmul(other, self)
+        return _ufuncs.matmul(other, self)
 
     def __imatmul__(self, other):
-        return _binary_ufuncs.matmul(self, other, out=self)
+        return _ufuncs.matmul(self, other, out=self)
 
     # unary ops
-    __invert__ = _unary_ufuncs.invert
-    __abs__ = _unary_ufuncs.absolute
-    __pos__ = _unary_ufuncs.positive
-    __neg__ = _unary_ufuncs.negative
+    __invert__ = _ufuncs.invert
+    __abs__ = _ufuncs.absolute
+    __pos__ = _ufuncs.positive
+    __neg__ = _ufuncs.negative
 
-    conjugate = _unary_ufuncs.conjugate
+    conjugate = _ufuncs.conjugate
     conj = conjugate
 
     ### methods to match namespace functions

torch_np/_binary_ufuncs.py renamed to torch_np/_ufuncs.py

@@ -2,13 +2,14 @@
 
 import torch
 
-from . import _helpers
-from ._detail import _binary_ufuncs
+from . import _binary_ufuncs_impl, _helpers, _unary_ufuncs_impl
 from ._normalizations import ArrayLike, DTypeLike, NDArray, SubokLike, normalizer
 
-__all__ = [
+# ############# Binary ufuncs ######################
+
+__binary__ = [
     name
-    for name in dir(_binary_ufuncs)
+    for name in dir(_binary_ufuncs_impl)
     if not name.startswith("_") and name not in ["torch", "matmul"]
 ]
 
@@ -76,7 +77,7 @@ def matmul(
         raise NotImplementedError
 
     # NB: do not broadcast input tensors against the out=... array
-    result = _binary_ufuncs.matmul(*tensors)
+    result = _binary_ufuncs_impl.matmul(*tensors)
     return result
 
 
@@ -110,7 +111,7 @@ def divmod(
         (x1, x2), out, True, casting, order, dtype, subok, signature, extobj
     )
 
-    result = _binary_ufuncs.divmod(*tensors)
+    result = _binary_ufuncs_impl.divmod(*tensors)
 
     return quot, rem
 
@@ -119,8 +120,8 @@ def divmod(
 # 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)
+for name in __binary__:
+    ufunc = getattr(_binary_ufuncs_impl, name)
     decorated = normalizer(deco_binary_ufunc(ufunc))
 
     decorated.__qualname__ = name  # XXX: is this really correct?
@@ -133,4 +134,64 @@ def modf(x, /, *args, **kwds):
     return rem, quot
 
 
-__all__ = __all__ + ["divmod", "modf", "matmul"]
+__binary__ = __binary__ + ["divmod", "modf", "matmul"]
+
+
+# ############# Unary ufuncs ######################
+
+
+__unary__ = [
+    name
+    for name in dir(_unary_ufuncs_impl)
+    if not name.startswith("_") and name != "torch"
+]
+
+
+def deco_unary_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(
+        x: 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(
+            (x,), out, where, casting, order, dtype, subok, signature, extobj
+        )
+        # now broadcast the input tensor against the out=... array
+        if out is not None:
+            # XXX: need to filter out noop broadcasts if t.shape == out.shape?
+            shape = out.shape
+            tensors = tuple(torch.broadcast_to(t, shape) for t in tensors)
+        result = torch_func(*tensors)
+        return result
+
+    return wrapped
+
+
+#
+# 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 __unary__:
+    ufunc = getattr(_unary_ufuncs_impl, name)
+    decorated = normalizer(deco_unary_ufunc(ufunc))
+
+    decorated.__qualname__ = name  # XXX: is this really correct?
+    decorated.__name__ = name
+    vars()[name] = decorated
+
+
+__all__ = __binary__ + __unary__

torch_np/_unary_ufuncs.py

@@ -1,5 +1,5 @@
 """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.
+This listing is further exported to public symbols in the `torch_np/_ufuncs.py` module.
 """
 
 import torch