dtypes and array scalars

torch_np/__init__.py

@@ -1,4 +1,5 @@
 from ._dtypes import *
+from ._scalar_types import *
 from ._wrapper import *
 from . import testing
 

torch_np/_dtypes.py

@@ -8,6 +8,12 @@
 import builtins
 import torch
 
+from . import _scalar_types
+
+
+__all__ = ['dtype_from_torch', 'dtype', 'typecodes', 'issubdtype']
+
+
 # Define analogs of numpy dtypes supported by pytorch.
 
 class dtype:
@@ -22,6 +28,9 @@ def __init__(self, name):
             _name = typecode_chars_dict[name]
         elif name in dt_aliases_dict:
             _name = dt_aliases_dict[name]
+       # the check must come last, so that 'name' is not a string
+        elif issubclass(name, _scalar_types.generic):
+            _name = name.name
         else:
             raise TypeError(f"data type '{name}' not understood")
         self._name = _name
@@ -30,6 +39,10 @@ def __init__(self, name):
     def name(self):
         return self._name
 
+    @property
+    def type(self):
+        return _scalar_types._typemap[self._name]
+
     @property
     def typecode(self):
         return _typecodes_from_dtype_dict[self._name]
@@ -104,21 +117,6 @@ def __repr__(self):
 }
 
 
-float16 = dtype("float16")
-float32 = dtype("float32")
-float64 = dtype("float64")
-complex64 = dtype("complex64")
-complex128 = dtype("complex128")
-uint8 = dtype("uint8")
-int8 = dtype("int8")
-int16 = dtype("int16")
-int32 = dtype("int32")
-int64 = dtype("int64")
-bool = dtype("bool")
-
-intp = int64   # XXX
-int_ = int64
-
 # Map the torch-suppored subset dtypes to local analogs
 # "quantized" types not available in numpy, skip
 _dtype_from_torch_dict = {
@@ -183,6 +181,23 @@ def is_integer(dtyp):
     return dtyp.typecode in typecodes['AllInteger']
 
 
+
+def issubclass_(arg, klass):
+    try:
+        return issubclass(arg, klass)
+    except TypeError:
+        return False
+
+
+def issubdtype(arg1, arg2):
+    # cf https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numerictypes.py#L356-L420
+    if not issubclass_(arg1, _scalar_types.generic):
+        arg1 = dtype(arg1).type
+    if not issubclass_(arg2, _scalar_types.generic):
+        arg2 = dtype(arg2).type
+    return issubclass(arg1, arg2)
+
+
 # The casting below is defined *with dtypes only*, so no value-based casting!
 
 # These two dicts are autogenerated with autogen/gen_dtypes.py,
@@ -216,6 +231,3 @@ def is_integer(dtyp):
 
 ########################## end autogenerated part
 
-
-__all__ = ['dtype_from_torch', 'dtype', 'typecodes'] + dt_names + ['intp', 'int_']
-

torch_np/_helpers.py

@@ -125,3 +125,9 @@ def allow_only_single_axis(axis):
     if len(axis) != 1:
         raise NotImplementedError("does not handle tuple axis")
     return axis[0]
+
+
+def to_tensors(*inputs):
+    """Convert all ndarrays from `inputs` to tensors."""
+    return tuple([value.get() if isinstance(value, ndarray) else value
+            for value in inputs])

torch_np/_scalar_types.py

@@ -0,0 +1,137 @@
+"""Replicate the NumPy scalar type hierarchy
+"""
+
+import abc
+import torch
+
+class generic(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def name(self):
+        return self.__class__.__name__
+
+    def __new__(self, value):
+        #
+        # Yes, a call to np.float32(4) produces a zero-dim array.
+        #
+        from . import _dtypes
+        from . import _ndarray
+
+        torch_dtype = _dtypes.torch_dtype_from(self.name)
+        tensor = torch.as_tensor(value, dtype=torch_dtype)
+        return _ndarray.ndarray._from_tensor_and_base(tensor, None)
+
+
+##### these are abstract types
+
+class number(generic):
+    pass
+
+
+class integer(generic):
+    pass
+
+
+class inexact(generic):
+    pass
+
+
+class signedinteger(generic):
+    pass
+
+
+class unsignedinteger(generic):
+    pass
+
+
+class inexact(generic):
+    pass
+
+
+class floating(generic):
+    pass
+
+
+class complexfloating(generic):
+    pass
+
+
+# ##### concrete types
+
+# signed integers
+
+class int8(signedinteger):
+    name = 'int8'
+
+
+class int16(signedinteger):
+    name = 'int16'
+
+
+class int32(signedinteger):
+    name = 'int32'
+
+
+class int64(signedinteger):
+    name = 'int64'
+
+
+# unsigned integers
+
+class uint8(unsignedinteger):
+    name = 'uint8'
+
+
+# floating point
+
+class float16(floating):
+    name = 'float16'
+
+
+class float32(floating):
+    name = 'float32'
+
+
+class float64(floating):
+    name = 'float64'
+
+
+class complex64(complexfloating):
+    name = 'complex64'
+
+
+class complex128(complexfloating):
+    name = 'complex128'
+
+
+class bool_(generic):
+    name = 'bool'
+
+
+# name aliases : FIXME (OS, bitness)
+intp = int64
+int_ = int64
+
+
+_typemap ={
+    'int8' : int8,
+    'int16' : int16,
+    'int32' : int32,
+    'int64' : int64,
+    'uint8' : uint8,
+    'float16': float16,
+    'float32': float32,
+    'float64': float64,
+    'complex64': complex64,
+    'complex128': complex128,
+    'bool': bool_
+}
+
+
+__all__ = list(_typemap.keys())
+__all__.remove('bool')
+
+__all__ += ['bool_', 'intp', 'int_']
+__all__ += ['generic', 'number',
+            'signedinteger', 'unsignedinteger',
+            'inexact', 'floating', 'complexfloating']

torch_np/_wrapper.py

@@ -193,8 +193,9 @@ def arange(start=None, stop=None, step=1, dtype=None, *, like=None):
     if dtype is None:
         dtype = _dtypes.default_int_type()
         dtype = result_type(start, stop, step, dtype)
-
     torch_dtype = _dtypes.torch_dtype_from(dtype)
+    start, stop, step = _helpers.to_tensors(start, stop, step)
+
     try:
         return asarray(torch.arange(start, stop, step, dtype=torch_dtype))
     except RuntimeError:

torch_np/tests/test_function_base.py

@@ -96,16 +96,15 @@ def test_arange_booleans(self):
         with pytest.raises(TypeError):
             np.arange(3, dtype="bool")
 
-    @pytest.mark.skip(reason='XXX: python scalars from array scalars')
     @pytest.mark.parametrize("which", [0, 1, 2])
     def test_error_paths_and_promotion(self, which):
-        args = [0, 1, 2]  # start, stop, and step
+        args = [0, 10, 2]  # start, stop, and step
         args[which] = np.float64(2.)  # should ensure float64 output
         assert np.arange(*args).dtype == np.float64
 
         # Cover stranger error path, test only to achieve code coverage!
         args[which] = [None, []]
-        with pytest.raises(ValueError):
+        with pytest.raises((ValueError, RuntimeError)):
             # Fails discovering start dtype
             np.arange(*args)
 

torch_np/tests/test_shape_base.py

@@ -464,7 +464,7 @@ def test_stack():
 
     # 0d input
     for input_ in [(1, 2, 3),
-###                   [np.int32(1), np.int32(2), np.int32(3)],    # XXX: numpy scalars?
+                   [np.int32(1), np.int32(2), np.int32(3)],
                    [np.array(1), np.array(2), np.array(3)]]:
         assert_array_equal(stack(input_), [1, 2, 3])
     # 1d input examples