Merge pull request #8 from Quansight-Labs/reductions

Implement reduction operations

Author: ev-br

autogen/numpy_api_dump.py

@@ -58,25 +58,13 @@ def seterrobj(errobj):
 
 
 
-def alen(a):
-    raise NotImplementedError
-
 
 def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
     raise NotImplementedError
 
 def alltrue(*args, **kwargs):
     raise NotImplementedError
 
-def amax(a, axis=None, out=None, keepdims=NoValue, initial=NoValue, where=NoValue):
-    raise NotImplementedError
-
-def amin(a, axis=None, out=None, keepdims=NoValue, initial=NoValue, where=NoValue):
-    raise NotImplementedError
-
-
-
-
 def append(arr, values, axis=None):
     raise NotImplementedError
 
@@ -95,10 +83,6 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
 def argsort(a, axis=-1, kind=None, order=None):
     raise NotImplementedError
 
-def argwhere(a):
-    raise NotImplementedError
-
-
 def array2string(a, max_line_width=None, precision=None, suppress_small=None, separator=' ', prefix='', style=NoValue, formatter=None, threshold=None, edgeitems=None, sign=None, floatmode=None, suffix='', *, legacy=None):
     raise NotImplementedError
 
@@ -126,9 +110,6 @@ def asfarray(a, dtype='numpy.float64'):
 def asmatrix(data, dtype=None):
     raise NotImplementedError
 
-def asscalar(a):
-    raise NotImplementedError
-
 
 def average(a, axis=None, weights=None, returned=False, *, keepdims=NoValue):
     raise NotImplementedError
@@ -189,8 +170,6 @@ def copyto(dst, src, casting='same_kind', where=True):
 def correlate(a, v, mode='valid'):
     raise NotImplementedError
 
-def count_nonzero(a, axis=None, *, keepdims=False):
-    raise NotImplementedError
 
 def cov(m, y=None, rowvar=True, bias=False, ddof=None, fweights=None, aweights=None, *, dtype=None):
     raise NotImplementedError
@@ -277,9 +256,6 @@ def find_common_type(array_types, scalar_types):
 def fix(x, out=None):
     raise NotImplementedError
 
-def flatnonzero(a):
-    raise NotImplementedError
-
 def flip(m, axis=None):
     raise NotImplementedError
 
@@ -444,8 +420,7 @@ def mask_indices(n, mask_func, k=0):
 def asmatrix(data, dtype=None):
     raise NotImplementedError
 
-def amax(a, axis=None, out=None, keepdims=NoValue, initial=NoValue, where=NoValue):
-    raise NotImplementedError
+
 
 def maximum_sctype(t):
     raise NotImplementedError
@@ -461,8 +436,7 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
 def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
     raise NotImplementedError
 
-def amin(a, axis=None, out=None, keepdims=NoValue, initial=NoValue, where=NoValue):
-    raise NotImplementedError
+
 
 def min_scalar_type(a, /):
     raise NotImplementedError

torch_np/__init__.py

@@ -6,3 +6,7 @@
 from ._binary_ufuncs import *
 from ._ndarray import can_cast, result_type, newaxis
 from ._util import AxisError
+
+
+inf = float('inf')
+nan = float('nan')

torch_np/_dtypes.py

@@ -18,6 +18,8 @@ def __init__(self, name):
             _name = python_types_dict[name]
         elif name in dt_names:
             _name = name
+        elif name in typecode_chars_dict:
+            _name = typecode_chars_dict[name]
         elif name in dt_aliases_dict:
             _name = dt_aliases_dict[name]
         else:
@@ -28,6 +30,10 @@ def __init__(self, name):
     def name(self):
         return self._name
 
+    @property
+    def typecode(self):
+        return _typecodes_from_dtype_dict[self._name]
+
     def __eq__(self, other):
         if isinstance(other, dtype):
             return self._name == other.name
@@ -62,6 +68,7 @@ def __repr__(self):
     'f8' : 'float64',
     'c8' : 'complex64',
     'c16': 'complex128',
+    '?'  : 'bool',
 }
 
 
@@ -72,6 +79,31 @@ def __repr__(self):
 }
 
 
+typecode_chars_dict = {
+    'e': 'float16',
+    'f': 'float32',
+    'd': 'float64',
+    'F': 'complex64',
+    'D': 'complex128',
+    'B': 'uint8',
+    'b': 'int8',
+    'h': 'int16',
+    'i': 'int32',
+    'l': 'int64',
+    '?': 'bool'
+}
+
+# reverse mapping
+_typecodes_from_dtype_dict = {typecode_chars_dict[key]: key
+                                for key in typecode_chars_dict}
+
+
+typecodes = {'All': 'efdFDBbhil?',
+        'AllFloat': 'efdFD',
+        'AllInteger': 'Bbhil',
+}
+
+
 float16 = dtype("float16")
 float32 = dtype("float32")
 float64 = dtype("float64")
@@ -84,6 +116,8 @@ def __repr__(self):
 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
@@ -132,6 +166,22 @@ def torch_dtype_from(dtyp):
         raise TypeError
 
 
+def default_int_type():
+    return dtype('int64')
+
+
+def default_float_type():
+    return dtype('float64')
+
+
+def is_floating(dtyp):
+    dtyp = dtype(dtyp)
+    return dtyp.typecode in typecodes['AllFloat']
+
+def is_integer(dtyp):
+    dtyp = dtype(dtyp)
+    return dtyp.typecode in typecodes['AllInteger']
+
 
 # The casting below is defined *with dtypes only*, so no value-based casting!
 
@@ -167,5 +217,5 @@ def torch_dtype_from(dtyp):
 ########################## end autogenerated part
 
 
-__all__ = ['dtype_from_torch', 'dtype'] + dt_names
+__all__ = ['dtype_from_torch', 'dtype', 'typecodes'] + dt_names + ['intp', 'int_']
 

torch_np/_helpers.py

@@ -1,6 +1,9 @@
+import operator
+
 import torch
 from . import _dtypes
 from ._ndarray import can_cast, ndarray, asarray
+from . import _util
 
 def cast_and_broadcast(arrays, out, casting):
     """Cast dtypes of arrays to out.dtype and broadcast if needed.
@@ -82,14 +85,43 @@ def axis_none_ravel(*arrays, axis=None):
         return arrays, axis
 
 
-def result_or_out(result, out=None):
+def result_or_out(result_tensor, out_array=None):
     """A helper for returns with out= argument."""
-    if out is not None:
-        if result.shape != out.shape:
+    if out_array is not None:
+        if result_tensor.shape != out_array.shape:
             raise ValueError
-        out_tensor = out.get()
-        out_tensor.copy_(result)
-        return out
+        out_tensor = out_array.get()
+        out_tensor.copy_(result_tensor)
+        return out_array
+    else:
+        return asarray(result_tensor)
+
+
+def apply_keepdims(tensor, axis, ndim):
+    if axis is None:
+        # tensor was a scalar
+        tensor = torch.full((1,)*ndim, fill_value=tensor)
     else:
-        return asarray(result)
+        shape = _util.expand_shape(tensor.shape, axis)
+        tensor = tensor.reshape(shape)
+    return tensor
 
+
+def standardize_axis_arg(axis, ndim):
+    """Return axis as either None or a tuple of normalized axes."""
+    if isinstance(axis, ndarray):
+        axis = operator.index(axis)
+
+    if axis is not None:
+        if not isinstance(axis, (list, tuple)):
+            axis = (axis,)
+        axis = _util.normalize_axis_tuple(axis, ndim)
+    return axis
+
+
+def allow_only_single_axis(axis):
+    if axis is None:
+        return axis
+    if len(axis) != 1:
+        raise NotImplementedError("does not handle tuple axis")
+    return axis[0]