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

Reviewed at #84

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -50,3 +50,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"]

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
 

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

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: