Merge pull request #91 from Quansight-Labs/rebase_stack

add divmod, matmul; add dtype=... arg to ufuncs, add several missing alisases.

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -1,11 +1,15 @@
 from typing import Optional
 
+import torch
+
 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"
+    name
+    for name in dir(_binary_ufuncs)
+    if not name.startswith("_") and name not in ["torch", "matmul"]
 ]
 
 
@@ -33,12 +37,49 @@ def wrapped(
         tensors = _helpers.ufunc_preprocess(
             (x1, x2), out, where, casting, order, dtype, subok, signature, extobj
         )
+        # now broadcast input tensors 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 _helpers.result_or_out(result, out)
 
     return wrapped
 
 
+#
+# matmul is special in that its `out=...` array does not broadcast x1 and x2.
+# E.g. consider x1.shape = (5, 2) and x2.shape = (2, 3). Then `out.shape` is (5, 3).
+#
+@normalizer
+def matmul(
+    x1: ArrayLike,
+    x2: ArrayLike,
+    /,
+    out: Optional[NDArray] = None,
+    *,
+    casting="same_kind",
+    order="K",
+    dtype: DTypeLike = None,
+    subok: SubokLike = False,
+    signature=None,
+    extobj=None,
+    axes=None,
+    axis=None,
+):
+    tensors = _helpers.ufunc_preprocess(
+        (x1, x2), out, True, casting, order, dtype, subok, signature, extobj
+    )
+    if axis is not None or axes is not None:
+        raise NotImplementedError
+
+    # NB: do not broadcast input tensors against the out=... array
+    result = _binary_ufuncs.matmul(*tensors)
+    return _helpers.result_or_out(result, out)
+
+
 #
 # 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
@@ -50,3 +91,58 @@ def wrapped(
     decorated.__qualname__ = name  # XXX: is this really correct?
     decorated.__name__ = name
     vars()[name] = decorated
+
+
+# a stub implementation of divmod, should be improved after
+# https://github.com/pytorch/pytorch/issues/90820 is fixed in pytorch
+#
+# Implementation details: we just call two ufuncs which have been created
+# just above, for x1 // x2 and x1 % x2.
+# This means we are normalizing x1, x2 in each of the ufuncs --- note that there
+# is no @normalizer on divmod.
+
+
+def divmod(
+    x1,
+    x2,
+    /,
+    out=None,
+    *,
+    where=True,
+    casting="same_kind",
+    order="K",
+    dtype=None,
+    subok: SubokLike = False,
+    signature=None,
+    extobj=None,
+):
+    out1, out2 = None, None
+    if out is not None:
+        out1, out2 = out
+
+    kwds = dict(
+        where=where,
+        casting=casting,
+        order=order,
+        dtype=dtype,
+        subok=subok,
+        signature=signature,
+        extobj=extobj,
+    )
+
+    # NB: use local names for
+    quot = floor_divide(x1, x2, out=out1, **kwds)
+    rem = remainder(x1, x2, out=out2, **kwds)
+
+    quot = _helpers.result_or_out(quot.tensor, out1)
+    rem = _helpers.result_or_out(rem.tensor, out2)
+
+    return quot, rem
+
+
+def modf(x, /, *args, **kwds):
+    quot, rem = divmod(x, 1, *args, **kwds)
+    return rem, quot
+
+
+__all__ = __all__ + ["divmod", "modf", "matmul"]

torch_np/_detail/_binary_ufuncs.py

@@ -38,16 +38,35 @@
     not_equal,
 )
 from torch import pow as power
-from torch import remainder, subtract
+from torch import remainder
+from torch import remainder as mod
+from torch import subtract, true_divide
 
 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
+    # work around:
+    #  - RuntimeError: expected scalar type Int but found Double
+    #  - RuntimeError: "addmm_impl_cpu_" not implemented for 'Bool'
+    #  - RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
     dtype = _dtypes_impl.result_type_impl((x.dtype, y.dtype))
-    x = _util.cast_if_needed(x, dtype)
-    y = _util.cast_if_needed(y, dtype)
+    is_bool = dtype == torch.bool
+    is_half = dtype == torch.float16
+
+    work_dtype = dtype
+    if is_bool:
+        work_dtype = torch.uint8
+    if is_half:
+        work_dtype = torch.float32
+
+    x = _util.cast_if_needed(x, work_dtype)
+    y = _util.cast_if_needed(y, work_dtype)
+
     result = torch.matmul(x, y)
+
+    if work_dtype != dtype:
+        result = result.to(dtype)
+
     return result

torch_np/_detail/_unary_ufuncs.py

@@ -7,6 +7,7 @@
 # renames
 from torch import absolute as fabs
 from torch import arccos, arccosh, arcsin, arcsinh, arctan, arctanh
+from torch import bitwise_not
 from torch import bitwise_not as invert
 from torch import ceil
 from torch import conj_physical as conjugate
@@ -31,6 +32,7 @@
 from torch import rad2deg
 from torch import rad2deg as degrees
 from torch import reciprocal
+from torch import round as fix
 from torch import round as rint
 from torch import sign, signbit, sin, sinh, sqrt, square, tan, tanh, trunc
 

torch_np/_detail/_util.py

@@ -137,7 +137,7 @@ def axis_none_ravel(*tensors, axis=None):
         return tensors, axis
 
 
-def cast_dont_broadcast(tensors, target_dtype, casting):
+def typecast_tensors(tensors, target_dtype, casting):
     """Dtype-cast tensors to target_dtype.
 
     Parameters
@@ -170,52 +170,6 @@ def cast_dont_broadcast(tensors, target_dtype, casting):
     return tuple(cast_tensors)
 
 
-def cast_and_broadcast(tensors, out_param, casting):
-    """
-    Parameters
-    ----------
-    tensors : iterable
-        tuple or list of torch.Tensors to broadcast/typecast
-    target_dtype : a torch.dtype object
-        The torch dtype to cast all tensors to
-    target_shape : tuple
-        The tensor shape to broadcast all `tensors` to
-    casting : str
-        The casting mode, see `np.can_cast`
-
-    Returns
-    -------
-    a tuple of torch.Tensors with dtype being the PyTorch counterpart
-    of the `target_dtype` and `target_shape`
-    """
-    if out_param is None:
-        return tensors
-
-    target_dtype, target_shape = out_param
-
-    can_cast = _dtypes_impl.can_cast_impl
-
-    processed_tensors = []
-    for tensor in tensors:
-        # check dtypes of x and out
-        if not can_cast(tensor.dtype, target_dtype, casting=casting):
-            raise TypeError(
-                f"Cannot cast array data from {tensor.dtype} to"
-                f" {target_dtype} according to the rule '{casting}'"
-            )
-
-        # cast arr if needed
-        tensor = cast_if_needed(tensor, target_dtype)
-
-        # `out` broadcasts `tensor`
-        if tensor.shape != target_shape:
-            tensor = torch.broadcast_to(tensor, target_shape)
-
-        processed_tensors.append(tensor)
-
-    return tuple(processed_tensors)
-
-
 def axis_expand_func(func, tensor, axis, *args, **kwds):
     """Generically handle axis arguments in reductions."""
     if axis is not None:

torch_np/_detail/implementations.py

@@ -308,7 +308,7 @@ def _concat_cast_helper(tensors, out=None, dtype=None, casting="same_kind"):
         out_dtype = _dtypes_impl.result_type_impl([t.dtype for t in tensors])
 
     # cast input arrays if necessary; do not broadcast them agains `out`
-    tensors = _util.cast_dont_broadcast(tensors, out_dtype, casting)
+    tensors = _util.typecast_tensors(tensors, out_dtype, casting)
 
     return tensors
 
@@ -467,7 +467,7 @@ def bincount(x, /, weights=None, minlength=0):
         x = x.new_empty(0, dtype=int)
 
     int_dtype = _dtypes_impl.default_int_dtype
-    (x,) = _util.cast_dont_broadcast((x,), int_dtype, casting="safe")
+    (x,) = _util.typecast_tensors((x,), int_dtype, casting="safe")
 
     result = torch.bincount(x, weights, minlength)
     return result

torch_np/_dtypes.py

@@ -174,7 +174,10 @@ class bool_(generic):
     "double": float64,
     "float_": float64,
     "csingle": complex64,
+    "singlecomplex": complex64,
     "cdouble": complex128,
+    "cfloat": complex128,
+    "complex_": complex128,
 }
 for name, obj in _name_aliases.items():
     globals()[name] = obj

torch_np/_helpers.py

@@ -1,58 +1,32 @@
 import torch
 
-from . import _dtypes
-from ._detail import _util
-
-
-def cast_and_broadcast(tensors, out, casting):
-    """Cast dtypes of arrays to out.dtype and broadcast if needed.
-
-    Parameters
-    ----------
-    arrays : sequence of arrays
-        Each element is broadcast against `out` and typecast to out.dtype
-    out : the "output" array
-        Not modified.
-    casting : str
-        One of numpy casting modes
-
-    Returns
-    -------
-    tensors : tuple of Tensors
-        Each tensor is dtype-cast and broadcast agains `out`, as needed
-
-    Notes
-    -----
-    The `out` arrays broadcasts and dtype-casts `arrays`, but not vice versa.
-
-    """
-    if out is None:
-        return tensors
-    else:
-        tensors = _util.cast_and_broadcast(
-            tensors, out.dtype.type.torch_dtype, out.shape, casting
-        )
-
-    return tuple(tensors)
+from ._detail import _dtypes_impl, _util
 
 
 def ufunc_preprocess(
     tensors, out, where, casting, order, dtype, subok, signature, extobj
 ):
+    """
+    Notes
+    -----
+    The `out` array broadcasts `tensors`, but not vice versa.
+    """
     # internal preprocessing or args in ufuncs (cf _unary_ufuncs, _binary_ufuncs)
     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.tensor.dtype, out.tensor.shape)
-
-    tensors = _util.cast_and_broadcast(tensors, out_shape_dtype, casting)
+    # dtype of the result: depends on both dtype=... and out=... arguments
+    if dtype is None:
+        out_dtype = None if out is None else out.dtype.torch_dtype
+    else:
+        out_dtype = (
+            dtype
+            if out is None
+            else _dtypes_impl.result_type_impl([dtype, out.dtype.torch_dtype])
+        )
 
+    if out_dtype:
+        tensors = _util.typecast_tensors(tensors, out_dtype, casting)
     return tensors
 
 

torch_np/_ndarray.py

@@ -266,6 +266,8 @@ def __rfloordiv__(self, other):
     def __ifloordiv__(self, other):
         return _binary_ufuncs.floor_divide(self, other, out=self)
 
+    __divmod__ = _binary_ufuncs.divmod
+
     # power, self**exponent
     __pow__ = __rpow__ = _binary_ufuncs.float_power
 
@@ -311,6 +313,14 @@ def __ilshift__(self, other):
     def __irshift__(self, other):
         return _binary_ufuncs.right_shift(self, other, out=self)
 
+    __matmul__ = _binary_ufuncs.matmul
+
+    def __rmatmul__(self, other):
+        return _binary_ufuncs.matmul(other, self)
+
+    def __imatmul__(self, other):
+        return _binary_ufuncs.matmul(self, other, out=self)
+
     # unary ops
     __invert__ = _unary_ufuncs.invert
     __abs__ = _unary_ufuncs.absolute

torch_np/_normalizations.py

@@ -45,7 +45,7 @@ def normalize_dtype(dtype, name=None):
     return torch_dtype
 
 
-def normalize_subok_like(arg, name):
+def normalize_subok_like(arg, name="subok"):
     if arg:
         raise ValueError(f"'{name}' parameter is not supported.")
 
@@ -88,7 +88,7 @@ def maybe_normalize(arg, parm, return_on_failure=_sentinel):
     """Normalize arg if a normalizer is registred."""
     normalizer = normalizers.get(parm.annotation, None)
     try:
-        return normalizer(arg) if normalizer else arg
+        return normalizer(arg, parm.name) if normalizer else arg
     except Exception as exc:
         if return_on_failure is not _sentinel:
             return return_on_failure