MAINT: finish up numeric binops/dunders

Author: ev-br

torch_np/_ndarray.py

@@ -182,7 +182,7 @@ def __len__(self):
 
     ### arithmetic ###
 
-    # add
+    # add, self + other
     def __add__(self, other):
         return _ufunc_impl.add(self, asarray(other))
 
@@ -193,7 +193,7 @@ def __iadd__(self, other):
         return _ufunc_impl.add(self, asarray(other), out=self)
 
 
-    # sub
+    # sub, self - other
     def __sub__(self, other):
         return _ufunc_impl.subtract(self, asarray(other))
 
@@ -204,7 +204,7 @@ def __isub__(self, other):
         return _ufunc_impl.subtract(self, asarray(other), out=self)
 
 
-    # mul
+    # mul, self * other
     def __mul__(self, other):
         return _ufunc_impl.multiply(self, asarray(other))
 
@@ -215,23 +215,40 @@ def __imul__(self, other):
         return _ufunc_impl.multiply(self, asarray(other), out=self)
 
 
+    # div, self / other
+    def __truediv__(self, other):
+        return _ufunc_impl.divide(self, asarray(other))
+
+    def __rtruediv__(self, other):
+        return _ufunc_impl.divide(self, asarray(other))
+
+    def __itruediv__(self, other):
+        return _ufunc_impl.divide(self, asarray(other), out=self)
+
 
+    # floordiv, self // other
     def __floordiv__(self, other):
-        other_tensor = asarray(other).get()
-        return asarray(self._tensor.__floordiv__(other_tensor))
+        return _ufunc_impl.floor_divide(self, asarray(other))
+
+    def __rfloordiv__(self, other):
+        return _ufunc_impl.floor_divide(self, asarray(other))
 
     def __ifloordiv__(self, other):
-        other_tensor = asarray(other).get()
-        return asarray(self._tensor.__ifloordiv__(other_tensor))
+        return _ufunc_impl.floor_divide(self, asarray(other), out=self)
 
-    def __truediv__(self, other):
-        other_tensor = asarray(other).get()
-        return asarray(self._tensor.__truediv__(other_tensor))
 
-    def __itruediv__(self, other):
-        other_tensor = asarray(other).get()
-        return asarray(self._tensor.__itruediv__(other_tensor))
+    # power, self**exponent
+    def __pow__(self, exponent):
+        return _ufunc_impl.float_power(self, asarray(exponent))
+
+    def __rpow__(self, exponent):
+        return _ufunc_impl.float_power(self, asarray(exponent))
 
+    def __ipow__(self, exponent):
+        return _ufunc_impl.float_power(self, asarray(exponent), out=self)
+
+
+    # FIXME ops and binops below
     def __mod__(self, other):
         other_tensor = asarray(other).get()
         return asarray(self._tensor.__mod__(other_tensor))
@@ -260,9 +277,7 @@ def __neg__(self):
         except RuntimeError as e:
             raise TypeError(e.args)
 
-    def __pow__(self, exponent):
-        exponent_tensor = asarray(exponent).get()
-        return asarray(self._tensor.__pow__(exponent_tensor))
+
 
     ### methods to match namespace functions
 

torch_np/tests/test_ufuncs_basic.py

@@ -93,23 +93,27 @@ def test_x_and_out_broadcast(self, ufunc):
     (np.add, operator.__add__, operator.__iadd__),
     (np.subtract, operator.__sub__, operator.__isub__),
     (np.multiply, operator.__mul__, operator.__imul__),
-# divide
-# true_divide?
+    (np.divide, operator.__truediv__, operator.__itruediv__),
+    (np.floor_divide, operator.__floordiv__, operator.__ifloordiv__),
+    (np.float_power, operator.__pow__, operator.__ipow__),
+
+
 # remainder vs fmod?
 # pow vs power vs float_power
 ]
 
 ufuncs_with_dunders = [ufunc for ufunc, _, _ in ufunc_op_iop_numeric]
+numeric_binary_ufuncs = [np.float_power, np.power,]
 
-numeric_binary_ufuncs = [np.float_power, np.power, 
 # these are not implemented for complex inputs
-#        np.hypot, np.arctan2, np.copysign,
-#        np.floor_divide, np.fmax, np.fmin, np.fmod,
-#        np.heaviside, np.logaddexp, np.logaddexp2, np.maximum, np.minimum,
+no_complex = [np.floor_divide, np.hypot, np.arctan2, np.copysign, np.fmax,
+        np.fmin, np.fmod, np.heaviside, np.logaddexp, np.logaddexp2,
+        np.maximum, np.minimum,
 ]
 
 parametrize_binary_ufuncs = pytest.mark.parametrize(
-        'ufunc', ufuncs_with_dunders + numeric_binary_ufuncs)
+        'ufunc', ufuncs_with_dunders + numeric_binary_ufuncs + no_complex)
+
 
 
 # TODO: these snowflakes need special handling
@@ -159,6 +163,9 @@ def test_xy_and_out_casting(self, ufunc, casting, out_dtype):
         x, y = self.get_xy(ufunc)
         out = np.empty_like(x, dtype=out_dtype)
 
+        if ufunc in no_complex and np.issubdtype(out_dtype, np.complexfloating):
+            pytest.skip(f'{ufunc} does not accept complex.')
+
         can_cast_x = np.can_cast(x, out_dtype, casting=casting)
         can_cast_y = np.can_cast(y, out_dtype, casting=casting)
 
@@ -208,12 +215,13 @@ def test_basic(self, ufunc, op, iop):
         assert_equal(op(a.tolist(), a), ufunc(a, a.tolist()))
 
         # __iadd__
-        a0 = np.array([1, 2, 3])
+        a0 = np.array([2, 4, 6])
         a = a0.copy()
+
         iop(a, 2)     # modifies a in-place
         assert_equal(a, op(a0, 2))
 
-        a0 = np.array([1, 2, 3])
+        a0 = np.array([2, 4, 6])
         a = a0.copy()
         iop(a, a)
         assert_equal(a, op(a0, a0))
@@ -225,6 +233,9 @@ def test_other_scalar(self, ufunc, op, iop, other_dtype):
         a = np.array([1, 2, 3])
         b = other_dtype(3)
 
+        if ufunc in no_complex and issubclass(other_dtype, np.complexfloating):
+            pytest.skip(f'{ufunc} does not accept complex.')
+
         # __op__
         result = op(a, b)
         assert_equal(result, ufunc(a, b))
@@ -253,9 +264,13 @@ def test_other_scalar(self, ufunc, op, iop, other_dtype):
     @pytest.mark.parametrize("other_dtype", dtypes_numeric)
     def test_other_array(self, ufunc, op, iop, other_dtype):
         """Test op/iop/rop when the other argument is an array of a different dtype."""
-        # __op__
         a = np.array([1, 2, 3])
         b = np.array([5, 6, 7], dtype=other_dtype)
+
+        if ufunc in no_complex and issubclass(other_dtype, np.complexfloating):
+            pytest.skip(f'{ufunc} does not accept complex.')
+
+        # __op__
         result = op(a, b)
         assert_equal(result, ufunc(a, b))
         assert result.dtype == np.result_type(a, b)