MAINT: rationalize the array/asarray split

Author: ev-br

torch_np/_dtypes.py

@@ -194,6 +194,8 @@ def torch_dtype_from(dtyp):
         raise TypeError
 
 
+# ### Defaults and dtype discovery
+
 def default_int_type():
     return dtype('int64')
 
@@ -202,14 +204,33 @@ def default_float_type():
     return dtype('float64')
 
 
+def default_complex_type():
+    return dtype('complex128')
+
+
 def is_floating(dtyp):
     dtyp = dtype(dtyp)
-    return dtyp.typecode in typecodes['AllFloat']
+    return issubclass(dtyp.type, _scalar_types.floating)
+
 
 def is_integer(dtyp):
     dtyp = dtype(dtyp)
-    return dtyp.typecode in typecodes['AllInteger']
-
+    return issubclass(dtyp.type, _scalar_types.integer)
+
+
+def get_default_dtype_for(dtyp):
+    typ = dtype(dtyp).type
+    if issubclass(typ, _scalar_types.integer):
+        result = default_int_type()
+    elif issubclass(typ, _scalar_types.floating):
+        result = default_float_type()
+    elif issubclass(typ, _scalar_types.complexfloating):
+        result = default_complex_type()
+    elif issubclass(typ, _scalar_types.bool_):
+        result = dtype('bool')
+    else:
+        raise TypeError("dtype %s not understood." % dtyp)
+    return result
 
 
 def issubclass_(arg, klass):

torch_np/_ndarray.py

@@ -159,6 +159,9 @@ def __hash__(self):
     def __float__(self):
         return float(self._tensor)
 
+    def __int__(self):
+        return int(self._tensor)
+
     # XXX : are single-element ndarrays scalars?
     def is_integer(self):
         if self.shape == ():
@@ -354,7 +357,7 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoVal
 
         if dtype is None:
             dtype = self.dtype
-        if not _dtypes.is_floating(dtype):
+        if _dtypes.is_integer(dtype):
             dtype = _dtypes.default_float_type()
         torch_dtype = _dtypes.torch_dtype_from(dtype)
 
@@ -374,7 +377,7 @@ def sum(self, axis=None, dtype=None, out=None, keepdims=NoValue,
 
         if dtype is None:
             dtype = self.dtype
-        if not _dtypes.is_floating(dtype):
+        if _dtypes.is_integer(dtype):
             dtype = _dtypes.default_float_type()
         torch_dtype = _dtypes.torch_dtype_from(dtype)
 
@@ -396,67 +399,80 @@ def __setitem__(self, index, value):
         return self._tensor.__setitem__(index, value)
 
 
-def asarray(a, dtype=None, order=None, *, like=None):
-    _util.subok_not_ok(like)
-    if order is not None:
+# This is the ideally the only place which talks to ndarray directly.
+# The rest goes through asarray (preferred) or array.
+
+def array(object, dtype=None, *, copy=True, order='K', subok=False, ndmin=0,
+          like=None):
+    _util.subok_not_ok(like, subok)
+    if order != 'K':
         raise NotImplementedError
 
-    if isinstance(a, ndarray):
-        if dtype is not None and dtype != a.dtype:
-            a = a.astype(dtype)
-        return a
+    # a happy path
+    if isinstance(object, ndarray):
+        if copy is False and dtype is None and ndmin <= object.ndim:
+            return object
 
-    if isinstance(a, (list, tuple)):
-        # handle lists of ndarrays, [1, [2, 3], ndarray(4)] etc
+    # lists of ndarrays: [1, [2, 3], ndarray(4)] convert to lists of lists
+    if isinstance(object, (list, tuple)):   
         a1 = []
-        for elem in a:
+        for elem in object:
             if isinstance(elem, ndarray):
                 a1.append(elem.get().tolist())
             else:
                 a1.append(elem)
+        object = a1
+
+    # get the tensor from "object"
+    if isinstance(object, ndarray):
+        tensor = object._tensor
+        base = object
+    elif isinstance(object, torch.Tensor):
+        tensor = object
+        base = None
     else:
-        a1 = a
+        tensor = torch.as_tensor(object)
+        base = None
 
-    torch_dtype = _dtypes.torch_dtype_from(dtype)
+        # At this point, `tensor.dtype` is the pytorch default. Our default may
+        # differ, so need to typecast. However, we cannot just do `tensor.to`,
+        # because if our desired dtype is wider then pytorch's, `tensor`
+        # may have lost precision:
 
-    # This and array(...) are the only places which talk to ndarray directly.
-    # The rest goes through asarray (preferred) or array.
-    out = ndarray()
-    tt = torch.as_tensor(a1, dtype=torch_dtype)
-    out._tensor = tt
-    return out
+        # int(torch.as_tensor(1e12)) - 1e12 equals -4096 (try it!)
 
+        # Therefore, we treat `tensor.dtype` as a hint, and convert the
+        # original object *again*, this time with an explicit dtype.
+        dtyp = _dtypes.dtype_from_torch(tensor.dtype)
+        default = _dtypes.get_default_dtype_for(dtyp)
+        torch_dtype = _dtypes.torch_dtype_from(default)
 
-def array(object, dtype=None, *, copy=True, order='K', subok=False, ndmin=0,
-          like=None):
-    _util.subok_not_ok(like, subok)
-    if order != 'K':
-        raise NotImplementedError
-
-    if isinstance(object, (list, tuple)):
-        obj = asarray(object)
-        return array(obj, dtype, copy=copy, order=order, subok=subok,
-                     ndmin=ndmin, like=like)
+        tensor = torch.as_tensor(object, dtype=torch_dtype)
 
-    if isinstance(object, ndarray):
-        result = object._tensor
-
-        if dtype != object.dtype:
-            torch_dtype = _dtypes.torch_dtype_from(dtype)
-            result = result.to(torch_dtype)
-    else:
+    # type cast if requested
+    if dtype is not None:
         torch_dtype = _dtypes.torch_dtype_from(dtype)
-        result = torch.as_tensor(object, dtype=torch_dtype)
+        tensor = tensor.to(torch_dtype)
+        base = None
 
+    # adjust ndim if needed
+    ndim_extra = ndmin - tensor.ndim
+    if ndim_extra > 0:
+        tensor = tensor.view((1,)*ndim_extra + tensor.shape)
+        base = None
+
+    # copy if requested
     if copy:
-        result = result.clone()    
+        tensor = tensor.clone()
+        base = None
 
-    ndim_extra = ndmin - result.ndim
-    if ndim_extra > 0:
-        result = result.reshape((1,)*ndim_extra + result.shape)
-    out = ndarray()
-    out._tensor = result
-    return out
+    return ndarray._from_tensor_and_base(tensor, base)
+
+
+def asarray(a, dtype=None, order=None, *, like=None):
+    if order is None:
+        order = 'K'
+    return array(a, dtype=dtype, order=order, like=like, copy=False, ndmin=0)
 
 
 

torch_np/tests/numpy_tests/core/test_scalarmath.py

@@ -191,7 +191,6 @@ def test_integers_to_negative_integer_power(self):
         # 1 ** -1 possible special case
         base = [np.array(1, dt)[()] for dt in 'bhilB']
         for i1, i2 in itertools.product(base, exp):
-            pass
             if i1.dtype != np.uint64:
                 assert_raises(ValueError, operator.pow, i1, i2)
             else:
@@ -435,15 +434,18 @@ def test_branches(self):
 
 class TestConversion:
     def test_int_from_long(self):
+        # NB: this test assumes that the default fp type is float64
         l = [1e6, 1e12, 1e18, -1e6, -1e12, -1e18]
         li = [10**6, 10**12, 10**18, -10**6, -10**12, -10**18]
         for T in [None, np.float64, np.int64]:
             a = np.array(l, dtype=T)
             assert_equal([int(_m) for _m in a], li)
 
+
     @pytest.mark.xfail(reason="pytorch does not emit this warning.")
     def test_iinfo_long_values_1(self):
         for code in 'bBh':
+
             with pytest.warns(DeprecationWarning):
                 res = np.array(np.iinfo(code).max + 1, dtype=code)
             tgt = np.iinfo(code).min

torch_np/tests/test_reductions.py

@@ -248,16 +248,19 @@ def test_mean_values(self):
         rmat = np.arange(20, dtype=float).reshape((4, 5))
         cmat = rmat + 1j*rmat
 
-        for mat in [rmat, cmat]:
-            for axis in [0, 1]:
-                tgt = mat.sum(axis=axis)
-                res = np.mean(mat, axis=axis) * mat.shape[axis]
-                assert_allclose(res, tgt)
-
-            for axis in [None]:
-                tgt = mat.sum(axis=axis)
-                res = np.mean(mat, axis=axis) * mat.size
-                assert_allclose(res, tgt)
+        import warnings
+        with warnings.catch_warnings():
+            warnings.simplefilter('error')
+            for mat in [rmat, cmat]:
+                for axis in [0, 1]:
+                    tgt = mat.sum(axis=axis)
+                    res = np.mean(mat, axis=axis) * mat.shape[axis]
+                    assert_allclose(res, tgt)
+
+                for axis in [None]:
+                    tgt = mat.sum(axis=axis)
+                    res = np.mean(mat, axis=axis) * mat.size
+                    assert_allclose(res, tgt)
 
     @pytest.mark.xfail(reason="see pytorch/gh-91597")
     def test_mean_float16(self):