MAINT: divmod + framework for out=(array, array)

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -4,7 +4,14 @@
 
 from . import _helpers
 from ._detail import _binary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, NDArray, OutArray, SubokLike, normalizer
+from ._normalizations import (
+    ArrayLike,
+    DTypeLike,
+    NDArray,
+    OutArray,
+    SubokLike,
+    normalizer,
+)
 
 __all__ = [
     name
@@ -93,53 +100,6 @@ def matmul(
     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

torch_np/_detail/_binary_ufuncs.py

@@ -70,3 +70,9 @@ def matmul(x, y):
         result = result.to(dtype)
 
     return result
+
+
+# a stub implementation of divmod, should be improved after
+# https://github.com/pytorch/pytorch/issues/90820 is fixed in pytorch
+def divmod(x, y):
+    return x // y, x % y

torch_np/_funcs.py

@@ -449,7 +449,7 @@ def cumprod(
 cumproduct = cumprod
 
 
-@normalizer
+@normalizer(promote_scalar_result=True)
 def quantile(
     a: ArrayLike,
     q: ArrayLike,

torch_np/_helpers.py

@@ -30,39 +30,6 @@ def ufunc_preprocess(
     return tensors
 
 
-# ### Return helpers: wrap a single tensor, a tuple of tensors, out= etc ###
-
-
-def result_or_out(result_tensor, out_array=None, promote_scalar=False):
-    """A helper for returns with out= argument.
-
-    If `promote_scalar is True`, then:
-        if result_tensor.numel() == 1 and out is zero-dimensional,
-            result_tensor is placed into the out array.
-    This weirdness is used e.g. in `np.percentile`
-    """
-    if out_array is not None:
-        if result_tensor.shape != out_array.shape:
-            can_fit = result_tensor.numel() == 1 and out_array.ndim == 0
-            if promote_scalar and can_fit:
-                result_tensor = result_tensor.squeeze()
-            else:
-                raise ValueError(
-                    f"Bad size of the out array: out.shape = {out_array.shape}"
-                    f" while result.shape = {result_tensor.shape}."
-                )
-        out_tensor = out_array.tensor
-        out_tensor.copy_(result_tensor)
-        return out_array
-    else:
-        from ._ndarray import ndarray
-
-        return ndarray(result_tensor)
-
-
-# ### Various ways of converting array-likes to tensors ###
-
-
 def ndarrays_to_tensors(*inputs):
     """Convert all ndarrays from `inputs` to tensors. (other things are intact)"""
     from ._ndarray import asarray, ndarray

torch_np/_normalizations.py

@@ -98,7 +98,6 @@ def maybe_normalize(arg, parm, return_on_failure=_sentinel):
             raise exc from None
 
 
-
 # ### Return value helpers ###
 
 
@@ -123,7 +122,6 @@ def copy_to_out(result_tensor, out_array, promote_scalar_result=False):
     return out_array
 
 
-
 def wrap_tensors(result, out, promote_scalar_result=False):
     from ._ndarray import ndarray
 
@@ -133,9 +131,11 @@ def wrap_tensors(result, out, promote_scalar_result=False):
         else:
             return copy_to_out(result, out, promote_scalar_result)
     elif isinstance(result, (tuple, list)):
-        return type(result)(
-            ndarray(x) if isinstance(x, torch.Tensor) else x for x in result
-        )
+        seq = type(result)
+        if out is None:
+            return seq(wrap_tensors(x, out) for x in result)
+        else:
+            return seq(wrap_tensors(x, out=o) for x, o in zip(result, out))
 
     return result
 
@@ -151,6 +151,7 @@ def array_or_scalar(values, py_type=float, return_scalar=False):
 
 # ### The main decorator to normalize arguments / postprocess the output ###
 
+
 def normalizer(_func=None, *, return_on_failure=_sentinel, promote_scalar_result=False):
     def normalizer_inner(func):
         @functools.wraps(func)

torch_np/_unary_ufuncs.py

@@ -4,7 +4,14 @@
 
 from . import _helpers
 from ._detail import _unary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, NDArray, OutArray, SubokLike, normalizer
+from ._normalizations import (
+    ArrayLike,
+    DTypeLike,
+    NDArray,
+    OutArray,
+    SubokLike,
+    normalizer,
+)
 
 __all__ = [
     name for name in dir(_unary_ufuncs) if not name.startswith("_") and name != "torch"

torch_np/_wrapper.py

@@ -12,7 +12,14 @@
 from . import _dtypes, _funcs, _helpers
 from ._detail import _dtypes_impl, _util
 from ._ndarray import asarray
-from ._normalizations import ArrayLike, DTypeLike, NDArray, SubokLike, OutArray, normalizer
+from ._normalizations import (
+    ArrayLike,
+    DTypeLike,
+    NDArray,
+    OutArray,
+    SubokLike,
+    normalizer,
+)
 
 NoValue = _util.NoValue