Refactor the internals to better separate wrappers from ops on tensors

autogen/gen_dtypes.py

@@ -22,11 +22,18 @@ def __init__(self, name):
          'int64',
          'bool']
 
+
+tmap = {dt: torch.as_tensor(np.ones(1, dtype=dt)).dtype for dt in dt_names}
+
+
+
 templ = """\
 {name} = dtype("{name}")
 """
 
 
+
+
 ############### Output the dtypes #############
 
 src_lines = [templ.format(name=name) for name in dt_names]
@@ -51,8 +58,8 @@ def generate_can_cast(casting):
         for dtyp2 in dt_names:
             can_cast = np.can_cast(np.dtype(dtyp1), np.dtype(dtyp2),
                                             casting=casting)
-            dct_dtyp1[dtyp2] = can_cast
-        dct[dtyp1] = dct_dtyp1
+            dct_dtyp1[tmap[dtyp2]] = can_cast
+        dct[tmap[dtyp1]] = dct_dtyp1
     return dct
 
 
@@ -63,8 +70,8 @@ def generate_result_type():
         dct_dtyp1 = {}
         for dtyp2 in dt_names:
             result_type = np.result_type(np.dtype(dtyp1), np.dtype(dtyp2))
-            dct_dtyp1[dtyp2] = result_type.name
-        dct[dtyp1] = dct_dtyp1
+            dct_dtyp1[tmap[dtyp2]] = tmap[result_type.name]
+        dct[tmap[dtyp1]] = dct_dtyp1
     return dct
 
 

torch_np/__init__.py

@@ -1,20 +1,15 @@
 from ._dtypes import *
-from ._scalar_types import *
+from ._detail._scalar_types import *
 from ._wrapper import *
 #from . import testing
 
 from ._unary_ufuncs import *
 from ._binary_ufuncs import *
 from ._ndarray import can_cast, result_type, newaxis
-from ._util import AxisError, UFuncTypeError
+from ._detail._util import AxisError, UFuncTypeError
 from ._getlimits import iinfo, finfo
 from ._getlimits import errstate
 
 inf = float('inf')
 nan = float('nan')
 
-
-#### HACK HACK HACK ####
-import torch
-torch.set_default_dtype(torch.float64)
-del torch

torch_np/_binary_ufuncs.py

@@ -1,79 +1,53 @@
-import functools
-import torch
+from ._decorators import deco_binary_ufunc_from_impl
+from ._detail import _ufunc_impl
 
-from . import _util
-
-from ._ndarray import asarray
-from . import _dtypes
-from . import _helpers
-
-from . import _ufunc_impl
-
-#
-# Functions in _ufunc_impl receive arrays, implement common tasks with ufunc args
-# and delegate heavy lifting to pytorch equivalents.
 #
 # Functions in this file implement binary ufuncs: wrap two first arguments in
 # asarray and delegate to functions from _ufunc_impl.
 #
-# One other user of _ufunc_impl functions in ndarray, where its __add__ method
-# calls _ufunc_impl.add and so on. Note that ndarray dunders already know
-# that its first arg is an array, so they only convert the second argument.
+# Functions in _detail/_ufunc_impl.py receive tensors, implement common tasks
+# with ufunc args, and delegate heavy lifting to pytorch equivalents.
 #
-# XXX: While it sounds tempting to merge _binary_ufuncs.py and _ufunc_impl.py
-# files, doing it would currently create import cycles.
-#
-
-# TODO: deduplicate with _unary_ufuncs/deco_unary_ufunc_from_impl,
-# _ndarray/asarray_replacer, and _wrapper/concatenate et al
-def deco_ufunc_from_impl(impl_func):
-    @functools.wraps(impl_func)
-    def wrapped(x1, x2, *args, **kwds):
-        x1_array = asarray(x1)
-        x2_array = asarray(x2)
-        return impl_func(x1_array, x2_array, *args, **kwds)
-    return wrapped
-
 
 # the list is autogenerated, cf autogen/gen_ufunc_2.py
-add = deco_ufunc_from_impl(_ufunc_impl.add)
-arctan2 = deco_ufunc_from_impl(_ufunc_impl.arctan2)
-bitwise_and = deco_ufunc_from_impl(_ufunc_impl.bitwise_and)
-bitwise_or = deco_ufunc_from_impl(_ufunc_impl.bitwise_or)
-bitwise_xor = deco_ufunc_from_impl(_ufunc_impl.bitwise_xor)
-copysign = deco_ufunc_from_impl(_ufunc_impl.copysign)
-divide = deco_ufunc_from_impl(_ufunc_impl.divide)
-equal = deco_ufunc_from_impl(_ufunc_impl.equal)
-float_power = deco_ufunc_from_impl(_ufunc_impl.float_power)
-floor_divide = deco_ufunc_from_impl(_ufunc_impl.floor_divide)
-fmax = deco_ufunc_from_impl(_ufunc_impl.fmax)
-fmin = deco_ufunc_from_impl(_ufunc_impl.fmin)
-fmod = deco_ufunc_from_impl(_ufunc_impl.fmod)
-gcd = deco_ufunc_from_impl(_ufunc_impl.gcd)
-greater = deco_ufunc_from_impl(_ufunc_impl.greater)
-greater_equal = deco_ufunc_from_impl(_ufunc_impl.greater_equal)
-heaviside = deco_ufunc_from_impl(_ufunc_impl.heaviside)
-hypot = deco_ufunc_from_impl(_ufunc_impl.hypot)
-lcm = deco_ufunc_from_impl(_ufunc_impl.lcm)
-ldexp = deco_ufunc_from_impl(_ufunc_impl.ldexp)
-left_shift = deco_ufunc_from_impl(_ufunc_impl.left_shift)
-less = deco_ufunc_from_impl(_ufunc_impl.less)
-less_equal = deco_ufunc_from_impl(_ufunc_impl.less_equal)
-logaddexp = deco_ufunc_from_impl(_ufunc_impl.logaddexp)
-logaddexp2 = deco_ufunc_from_impl(_ufunc_impl.logaddexp2)
-logical_and = deco_ufunc_from_impl(_ufunc_impl.logical_and)
-logical_or = deco_ufunc_from_impl(_ufunc_impl.logical_or)
-logical_xor = deco_ufunc_from_impl(_ufunc_impl.logical_xor)
-matmul = deco_ufunc_from_impl(_ufunc_impl.matmul)
-maximum = deco_ufunc_from_impl(_ufunc_impl.maximum)
-minimum = deco_ufunc_from_impl(_ufunc_impl.minimum)
-remainder = deco_ufunc_from_impl(_ufunc_impl.remainder)
-multiply = deco_ufunc_from_impl(_ufunc_impl.multiply)
-nextafter = deco_ufunc_from_impl(_ufunc_impl.nextafter)
-not_equal = deco_ufunc_from_impl(_ufunc_impl.not_equal)
-power = deco_ufunc_from_impl(_ufunc_impl.power)
-remainder = deco_ufunc_from_impl(_ufunc_impl.remainder)
-right_shift = deco_ufunc_from_impl(_ufunc_impl.right_shift)
-subtract = deco_ufunc_from_impl(_ufunc_impl.subtract)
-divide = deco_ufunc_from_impl(_ufunc_impl.divide)
+add = deco_binary_ufunc_from_impl(_ufunc_impl.add)
+arctan2 = deco_binary_ufunc_from_impl(_ufunc_impl.arctan2)
+bitwise_and = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_and)
+bitwise_or = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_or)
+bitwise_xor = deco_binary_ufunc_from_impl(_ufunc_impl.bitwise_xor)
+copysign = deco_binary_ufunc_from_impl(_ufunc_impl.copysign)
+divide = deco_binary_ufunc_from_impl(_ufunc_impl.divide)
+equal = deco_binary_ufunc_from_impl(_ufunc_impl.equal)
+float_power = deco_binary_ufunc_from_impl(_ufunc_impl.float_power)
+floor_divide = deco_binary_ufunc_from_impl(_ufunc_impl.floor_divide)
+fmax = deco_binary_ufunc_from_impl(_ufunc_impl.fmax)
+fmin = deco_binary_ufunc_from_impl(_ufunc_impl.fmin)
+fmod = deco_binary_ufunc_from_impl(_ufunc_impl.fmod)
+gcd = deco_binary_ufunc_from_impl(_ufunc_impl.gcd)
+greater = deco_binary_ufunc_from_impl(_ufunc_impl.greater)
+greater_equal = deco_binary_ufunc_from_impl(_ufunc_impl.greater_equal)
+heaviside = deco_binary_ufunc_from_impl(_ufunc_impl.heaviside)
+hypot = deco_binary_ufunc_from_impl(_ufunc_impl.hypot)
+lcm = deco_binary_ufunc_from_impl(_ufunc_impl.lcm)
+ldexp = deco_binary_ufunc_from_impl(_ufunc_impl.ldexp)
+left_shift = deco_binary_ufunc_from_impl(_ufunc_impl.left_shift)
+less = deco_binary_ufunc_from_impl(_ufunc_impl.less)
+less_equal = deco_binary_ufunc_from_impl(_ufunc_impl.less_equal)
+logaddexp = deco_binary_ufunc_from_impl(_ufunc_impl.logaddexp)
+logaddexp2 = deco_binary_ufunc_from_impl(_ufunc_impl.logaddexp2)
+logical_and = deco_binary_ufunc_from_impl(_ufunc_impl.logical_and)
+logical_or = deco_binary_ufunc_from_impl(_ufunc_impl.logical_or)
+logical_xor = deco_binary_ufunc_from_impl(_ufunc_impl.logical_xor)
+matmul = deco_binary_ufunc_from_impl(_ufunc_impl.matmul)
+maximum = deco_binary_ufunc_from_impl(_ufunc_impl.maximum)
+minimum = deco_binary_ufunc_from_impl(_ufunc_impl.minimum)
+remainder = deco_binary_ufunc_from_impl(_ufunc_impl.remainder)
+multiply = deco_binary_ufunc_from_impl(_ufunc_impl.multiply)
+nextafter = deco_binary_ufunc_from_impl(_ufunc_impl.nextafter)
+not_equal = deco_binary_ufunc_from_impl(_ufunc_impl.not_equal)
+power = deco_binary_ufunc_from_impl(_ufunc_impl.power)
+remainder = deco_binary_ufunc_from_impl(_ufunc_impl.remainder)
+right_shift = deco_binary_ufunc_from_impl(_ufunc_impl.right_shift)
+subtract = deco_binary_ufunc_from_impl(_ufunc_impl.subtract)
+divide = deco_binary_ufunc_from_impl(_ufunc_impl.divide)
 

torch_np/_decorators.py

@@ -0,0 +1,107 @@
+import functools
+import operator
+
+import torch
+
+from . import _dtypes
+from . import _helpers
+from ._detail import _util
+
+NoValue = None
+
+
+def dtype_to_torch(func):
+    @functools.wraps(func)
+    def wrapped(*args, dtype=None, **kwds):
+        torch_dtype = None
+        if dtype is not None:
+            dtype = _dtypes.dtype(dtype)
+            torch_dtype = dtype._scalar_type.torch_dtype
+        return func(*args, dtype=torch_dtype, **kwds)
+    return wrapped
+
+
+def emulate_out_arg(func):
+    """Simulate the out=... handling: move the result tensor to the out array.
+
+    With this decorator, the inner function just does not see the out array.
+    """
+    @functools.wraps(func)
+    def wrapped(*args, out=None, **kwds):
+        result_tensor = func(*args, **kwds)
+        return _helpers.result_or_out(result_tensor, out)
+
+    return wrapped
+
+
+def out_shape_dtype(func):
+    """Handle out=... kwarg for ufuncs.
+
+    With ufuncs, `out` array can typcast and broadcast ufunc arguments, hence
+    extract the shape and dtype of the tensor which backs the `out` array
+    and pass these through.
+    """
+    @functools.wraps(func)
+    def wrapped(*args, out=None, **kwds):
+        if out is not None:
+            kwds.update({'out_shape_dtype': (out.get().dtype, out.get().shape)})
+        result_tensor = func(*args, **kwds)
+        return _helpers.result_or_out(result_tensor, out)
+
+    return wrapped
+
+
+def deco_unary_ufunc_from_impl(impl_func):
+    @functools.wraps(impl_func)
+    @dtype_to_torch
+    @out_shape_dtype
+    def wrapped(x1, *args, **kwds):
+        from ._ndarray import asarray
+        x1_tensor = asarray(x1).get()
+        result = impl_func((x1_tensor,), *args, **kwds)
+        return result
+    return wrapped
+
+
+# TODO: deduplicate with _ndarray/asarray_replacer,
+# and _wrapper/concatenate et al
+def deco_binary_ufunc_from_impl(impl_func):
+    @functools.wraps(impl_func)
+    @dtype_to_torch
+    @out_shape_dtype
+    def wrapped(x1, x2, *args, **kwds):
+        from ._ndarray import asarray
+        x1_tensor = asarray(x1).get()
+        x2_tensor = asarray(x2).get()
+        return impl_func((x1_tensor, x2_tensor), *args, **kwds)
+    return wrapped
+
+
+
+def axis_keepdims_wrapper(func):
+    """`func` accepts an array-like as a 1st arg, returns a tensor.
+
+    This decorator implements the generic handling of axis, out and keepdims
+    arguments for reduction functions.
+
+    Note that we peel off `out=...` and `keepdims=...` args (torch functions never
+    see them). The `axis` argument we normalize and pass through to pytorch functions.
+
+    """
+    # XXX: move this out of _ndarray.py (circular imports)
+    #
+    # TODO: 1. get rid of _helpers.result_or_out
+    #       2. sort out function signatures: how they flow through all decorators etc
+    @functools.wraps(func)
+    def wrapped(a, axis=None, keepdims=NoValue, *args, **kwds):
+        from ._ndarray import ndarray, asarray
+        tensor = asarray(a).get()
+
+        # standardize the axis argument
+        if isinstance(axis, ndarray):
+            axis = operator.index(axis)
+
+        result = _util.axis_keepdims(func, tensor, axis, keepdims, *args, **kwds)
+        return result
+
+    return wrapped