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,39 +1,13 @@
-import functools
-import torch
+from ._decorators import deco_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)

torch_np/_decorators.py

@@ -0,0 +1,112 @@
+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  # XXX: why
+    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 *for functions which do not need it*.
+
+    With this decorator, the inner function just does not see the out array.
+    """
+    def wrapped(*args, out=None, **kwds):
+        from ._ndarray import ndarray
+        if out is not None:
+            if not isinstance(out, ndarray):
+                raise TypeError("Return arrays must be of ArrayType")        
+        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.
+    """
+    def wrapped(*args, out=None, **kwds):
+        from ._ndarray import ndarray
+        if out is not None:
+            if not isinstance(out, ndarray):
+                raise TypeError("Return arrays must be of ArrayType")
+            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_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, out=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