MAINT: expand types/dtypes, adapt test_scalarmath

Author: ev-br

torch_np/_dtypes.py

@@ -67,6 +67,7 @@ def __repr__(self):
 
     __str__ = __repr__
 
+    @property
     def itemsize(self):
         elem = self.type(1)
         return elem.get().element_size()

torch_np/_ndarray.py

@@ -81,6 +81,18 @@ def base(self):
     def T(self):
         return self.transpose()
 
+    @property
+    def real(self):
+        return asarray(self._tensor.real)
+
+    @property
+    def imag(self):
+        try:
+            return asarray(self._tensor.imag)
+        except RuntimeError:
+            zeros = torch.zeros_like(self._tensor)
+            return ndarray._from_tensor_and_base(zeros, None)
+
     # ctors
     def astype(self, dtype):
         newt = ndarray()
@@ -102,6 +114,13 @@ def __str__(self):
 
     ### comparisons ###
     def __eq__(self, other):
+        try:
+            t_other = asarray(other).get
+        except RuntimeError:
+            # Failed to convert other to array: definitely not equal.
+            # TODO: generalize, delegate to ufuncs
+            falsy = torch.full(self.shape, fill_value=False, dtype=bool)
+            return asarray(falsy)
         return asarray(self._tensor == asarray(other).get())
 
     def __neq__(self, other):
@@ -119,7 +138,6 @@ def __ge__(self, other):
     def __le__(self, other):
         return asarray(self._tensor <= asarray(other).get())
 
-
     def __bool__(self):
         try:
             return bool(self._tensor)
@@ -167,7 +185,10 @@ def __iadd__(self, other):
 
     def __sub__(self, other):
         other_tensor = asarray(other).get()
-        return asarray(self._tensor.__sub__(other_tensor))
+        try:
+            return asarray(self._tensor.__sub__(other_tensor))
+        except RuntimeError as e:
+            raise TypeError(e.args)
 
     def __mul__(self, other):
         other_tensor = asarray(other).get()
@@ -177,10 +198,30 @@ def __rmul__(self, other):
         other_tensor = asarray(other).get()
         return asarray(self._tensor.__rmul__(other_tensor))
 
+    def __floordiv__(self, other):
+        other_tensor = asarray(other).get()
+        return asarray(self._tensor.__floordiv__(other_tensor))
+
+    def __ifloordiv__(self, other):
+        other_tensor = asarray(other).get()
+        return asarray(self._tensor.__ifloordiv__(other_tensor))
+
     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))
+
+    def __mod__(self, other):
+        other_tensor = asarray(other).get()
+        return asarray(self._tensor.__mod__(other_tensor))
+
+    def __imod__(self, other):
+        other_tensor = asarray(other).get()
+        return asarray(self._tensor.__imod__(other_tensor))
+
     def __or__(self, other):
         other_tensor = asarray(other).get()
         return asarray(self._tensor.__or__(other_tensor))
@@ -189,10 +230,22 @@ def __ior__(self, other):
         other_tensor = asarray(other).get()
         return asarray(self._tensor.__ior__(other_tensor))
 
-
     def __invert__(self):
         return asarray(self._tensor.__invert__())
 
+    def __abs__(self):
+        return asarray(self._tensor.__abs__())
+
+    def __neg__(self):
+        try:
+            return asarray(self._tensor.__neg__())
+        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
 
     def squeeze(self, axis=None):

torch_np/_scalar_types.py

@@ -25,7 +25,12 @@ def __new__(self, value):
         if isinstance(value, _ndarray.ndarray):
             tensor = value.get()
         else:
-            tensor = torch.as_tensor(value, dtype=torch_dtype)
+            try:
+                tensor = torch.as_tensor(value, dtype=torch_dtype)
+            except RuntimeError as e:
+                if "Overflow" in str(e):
+                    raise OverflowError(e.args)
+                raise e
         #
         # With numpy:
         # >>> a = np.ones(3)
@@ -135,6 +140,7 @@ class bool_(generic):
 half = float16
 single = float32
 double = float64
+float_ = float64
 
 csingle = complex64
 cdouble = complex128
@@ -169,8 +175,8 @@ class bool_(generic):
 __all__ = list(_typemap.keys())
 __all__.remove('bool')
 
-__all__ += ['bool_', 'intp', 'int_', 'intc', 'byte', 'short', 'longlong', 'ubyte', 'half', 'single', 'double',
-'csingle', 'cdouble']
+__all__ += ['bool_', 'intp', 'int_', 'intc', 'byte', 'short', 'longlong',
+            'ubyte', 'half', 'single', 'double', 'csingle', 'cdouble', 'float_']
 __all__ += ['sctypes']
 __all__ += ['generic', 'number',
             'integer', 'signedinteger', 'unsignedinteger',

torch_np/_wrapper.py

@@ -510,6 +510,11 @@ def argwhere(a):
     return asarray(torch.argwhere(tensor))
 
 
+def abs(a):
+    # FIXME: should go the other way, together with other ufuncs
+    arr = asarray(a)
+    return a.__abs__()
+
 from ._ndarray import axis_out_keepdims_wrapper
 
 @axis_out_keepdims_wrapper
@@ -702,18 +707,14 @@ def angle(z, deg=False):
     return result
 
 
-@asarray_replacer()
 def real(a):
-    return torch.real(a)
+    arr = asarray(a)
+    return arr.real
 
 
-@asarray_replacer()
 def imag(a):
-    # torch.imag raises on real-valued inputs
-    if torch.is_complex(a):
-        return torch.imag(a) 
-    else:
-        return torch.zeros_like(a)
+    arr = asarray(a)
+    return arr.imag
 
 
 @asarray_replacer()

torch_np/testing/__init__.py

@@ -1,5 +1,6 @@
 from .utils import (assert_equal, assert_array_equal, assert_almost_equal,
         assert_warns, assert_)
+from .utils import _gen_alignment_data
 
 from .testing import assert_allclose    # FIXME
 

torch_np/testing/utils.py

@@ -869,9 +869,9 @@ def assert_array_almost_equal(x, y, decimal=6, err_msg='', verbose=True):
 
     """
     __tracebackhide__ = True  # Hide traceback for py.test
-    from numpy.core import number, float_, result_type, array
-    from numpy.core.numerictypes import issubdtype
-    from numpy.core.fromnumeric import any as npany
+    from torch_np import number, float_, result_type, array
+    from torch_np import issubdtype
+    from torch_np import any as npany
 
     def compare(x, y):
         try: