ENH: add a naive divmod, un-xfail relevant tests

torch_np/_binary_ufuncs.py

@@ -57,3 +57,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.get(), out1)  # FIXME: .get() -> .tensor
+    rem = _helpers.result_or_out(rem.get(), out2)
+
+    return quot, rem
+
+
+def modf(x, /, *args, **kwds):
+    quot, rem = divmod(x, 1, *args, **kwds)
+    return rem, quot
+
+
+__all__ = __all__ + ["divmod", "modf"]

torch_np/_ndarray.py

@@ -272,6 +272,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
 

torch_np/_normalizations.py

@@ -1,12 +1,14 @@
 """ "Normalize" arguments: convert array_likes to tensors, dtypes to torch dtypes and so on.
 """
+import functools
+import inspect
 import operator
 import typing
 from typing import Optional, Sequence, Union
 
 import torch
 
-from . import _helpers
+from . import _dtypes, _helpers
 
 ArrayLike = typing.TypeVar("ArrayLike")
 DTypeLike = typing.TypeVar("DTypeLike")
@@ -22,10 +24,6 @@
 NDArrayOrSequence = Union[NDArray, Sequence[NDArray]]
 OutArray = typing.TypeVar("OutArray")
 
-import inspect
-
-from . import _dtypes
-
 
 def normalize_array_like(x, name=None):
     (tensor,) = _helpers.to_tensors(x)
@@ -52,7 +50,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.")
 
@@ -87,7 +85,6 @@ def normalize_ndarray(arg, name=None):
     AxisLike: normalize_axis_like,
 }
 
-import functools
 
 _sentinel = object()
 
@@ -97,7 +94,7 @@ def normalize_this(arg, parm, return_on_failure=_sentinel):
     normalizer = normalizers.get(parm.annotation, None)
     if normalizer:
         try:
-            return normalizer(arg)
+            return normalizer(arg, parm.name)
         except Exception as exc:
             if return_on_failure is not _sentinel:
                 return return_on_failure
@@ -108,6 +105,44 @@ def normalize_this(arg, parm, return_on_failure=_sentinel):
         return arg
 
 
+# postprocess return values
+
+
+def postprocess_ndarray(result, **kwds):
+    return _helpers.array_from(result)
+
+
+def postprocess_out(result, **kwds):
+    result, out = result
+    return _helpers.result_or_out(result, out, **kwds)
+
+
+def postprocess_tuple(result, **kwds):
+    return _helpers.tuple_arrays_from(result)
+
+
+def postprocess_list(result, **kwds):
+    return list(_helpers.tuple_arrays_from(result))
+
+
+def postprocess_variadic(result, **kwds):
+    # a variadic return: a single NDArray or tuple/list of NDArrays, e.g. atleast_1d
+    if isinstance(result, (tuple, list)):
+        seq = type(result)
+        return seq(_helpers.tuple_arrays_from(result))
+    else:
+        return _helpers.array_from(result)
+
+
+postprocessors = {
+    NDArray: postprocess_ndarray,
+    OutArray: postprocess_out,
+    NDArrayOrSequence: postprocess_variadic,
+    tuple[NDArray]: postprocess_tuple,
+    list[NDArray]: postprocess_list,
+}
+
+
 def normalizer(_func=None, *, return_on_failure=_sentinel, promote_scalar_out=False):
     def normalizer_inner(func):
         @functools.wraps(func)
@@ -154,33 +189,17 @@ def wrapped(*args, **kwds):
                 raise TypeError(
                     f"{func.__name__}() takes {len(ba.args)} positional argument but {len(args)} were given."
                 )
+
             # finally, pass normalized arguments through
             result = func(*ba.args, **ba.kwargs)
 
             # handle returns
             r = sig.return_annotation
-            if r == NDArray:
-                return _helpers.array_from(result)
-            elif r == inspect._empty:
-                return result
-            elif hasattr(r, "__origin__") and r.__origin__ in (list, tuple):
-                # this is tuple[NDArray] or list[NDArray]
-                # XXX: change to separate tuple and list normalizers?
-                return r.__origin__(_helpers.tuple_arrays_from(result))
-            elif r == NDArrayOrSequence:
-                # a variadic return: a single NDArray or tuple/list of NDArrays, e.g. atleast_1d
-                if isinstance(result, (tuple, list)):
-                    seq = type(result)
-                    return seq(_helpers.tuple_arrays_from(result))
-                else:
-                    return _helpers.array_from(result)
-            elif r == OutArray:
-                result, out = result
-                return _helpers.result_or_out(
-                    result, out, promote_scalar=promote_scalar_out
-                )
-            else:
-                raise ValueError(f"Unknown return annotation {return_annotation}")
+            postprocess = postprocessors.get(r, None)
+            if postprocess:
+                kwds = {"promote_scalar": promote_scalar_out}
+                result = postprocess(result, **kwds)
+            return result
 
         return wrapped
 

torch_np/tests/numpy_tests/core/test_scalarmath.py

@@ -262,10 +262,6 @@ class TestModulus:
     def test_modulus_basic(self):
         dt = np.typecodes['AllInteger'] + np.typecodes['Float']
         for op in [floordiv_and_mod, divmod]:
-
-            if op == divmod:
-                pytest.xfail(reason="__divmod__ not implemented")
-
             for dt1, dt2 in itertools.product(dt, dt):
                 for sg1, sg2 in itertools.product(_signs(dt1), _signs(dt2)):
                     fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
@@ -279,7 +275,7 @@ def test_modulus_basic(self):
                     else:
                         assert_(b > rem >= 0, msg)
 
-    @pytest.mark.xfail(reason='divmod not implemented')
+    @pytest.mark.slow
     def test_float_modulus_exact(self):
         # test that float results are exact for small integers. This also
         # holds for the same integers scaled by powers of two.
@@ -311,10 +307,6 @@ def test_float_modulus_roundoff(self):
         # gh-6127
         dt = np.typecodes['Float']
         for op in [floordiv_and_mod, divmod]:
-
-            if op == divmod:
-                pytest.xfail(reason="__divmod__ not implemented")
-
             for dt1, dt2 in itertools.product(dt, dt):
                 for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
                     fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
@@ -329,41 +321,42 @@ def test_float_modulus_roundoff(self):
                     else:
                         assert_(b > rem >= 0, msg)
 
-    @pytest.mark.skip(reason='float16 on cpu is incomplete in pytorch')
-    def test_float_modulus_corner_cases(self):
-        # Check remainder magnitude.
-        for dt in np.typecodes['Float']:
-            b = np.array(1.0, dtype=dt)
-            a = np.nextafter(np.array(0.0, dtype=dt), -b)
-            rem = operator.mod(a, b)
-            assert_(rem <= b, 'dt: %s' % dt)
-            rem = operator.mod(-a, -b)
-            assert_(rem >= -b, 'dt: %s' % dt)
+    @pytest.mark.parametrize('dt', np.typecodes['Float'])
+    def test_float_modulus_corner_cases(self, dt):
+        if dt == 'e':
+            pytest.xfail(reason="RuntimeError: 'nextafter_cpu' not implemented for 'Half'")
+
+        b = np.array(1.0, dtype=dt)
+        a = np.nextafter(np.array(0.0, dtype=dt), -b)
+        rem = operator.mod(a, b)
+        assert_(rem <= b, 'dt: %s' % dt)
+        rem = operator.mod(-a, -b)
+        assert_(rem >= -b, 'dt: %s' % dt)
 
         # Check nans, inf
-        with suppress_warnings() as sup:
-            sup.filter(RuntimeWarning, "invalid value encountered in remainder")
-            sup.filter(RuntimeWarning, "divide by zero encountered in remainder")
-            sup.filter(RuntimeWarning, "divide by zero encountered in floor_divide")
-            sup.filter(RuntimeWarning, "divide by zero encountered in divmod")
-            sup.filter(RuntimeWarning, "invalid value encountered in divmod")
-            for dt in np.typecodes['Float']:
-                fone = np.array(1.0, dtype=dt)
-                fzer = np.array(0.0, dtype=dt)
-                finf = np.array(np.inf, dtype=dt)
-                fnan = np.array(np.nan, dtype=dt)
-                rem = operator.mod(fone, fzer)
-                assert_(np.isnan(rem), 'dt: %s' % dt)
-                # MSVC 2008 returns NaN here, so disable the check.
-                #rem = operator.mod(fone, finf)
-                #assert_(rem == fone, 'dt: %s' % dt)
-                rem = operator.mod(fone, fnan)
-                assert_(np.isnan(rem), 'dt: %s' % dt)
-                rem = operator.mod(finf, fone)
-                assert_(np.isnan(rem), 'dt: %s' % dt)
-                for op in [floordiv_and_mod, divmod]:
-                    div, mod = op(fone, fzer)
-                    assert_(np.isinf(div)) and assert_(np.isnan(mod))
+   #     with suppress_warnings() as sup:
+   #         sup.filter(RuntimeWarning, "invalid value encountered in remainder")
+   #         sup.filter(RuntimeWarning, "divide by zero encountered in remainder")
+   #         sup.filter(RuntimeWarning, "divide by zero encountered in floor_divide")
+   #         sup.filter(RuntimeWarning, "divide by zero encountered in divmod")
+   #         sup.filter(RuntimeWarning, "invalid value encountered in divmod")
+        for dt in np.typecodes['Float']:
+            fone = np.array(1.0, dtype=dt)
+            fzer = np.array(0.0, dtype=dt)
+            finf = np.array(np.inf, dtype=dt)
+            fnan = np.array(np.nan, dtype=dt)
+            rem = operator.mod(fone, fzer)
+            assert_(np.isnan(rem), 'dt: %s' % dt)
+            # MSVC 2008 returns NaN here, so disable the check.
+            #rem = operator.mod(fone, finf)
+            #assert_(rem == fone, 'dt: %s' % dt)
+            rem = operator.mod(fone, fnan)
+            assert_(np.isnan(rem), 'dt: %s' % dt)
+            rem = operator.mod(finf, fone)
+            assert_(np.isnan(rem), 'dt: %s' % dt)
+            for op in [floordiv_and_mod, divmod]:
+                div, mod = op(fone, fzer)
+                assert_(np.isinf(div)) and assert_(np.isnan(mod))
 
 
 class TestComplexDivision: