MAINT: address review comments

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -1,4 +1,4 @@
-from ._decorators import deco_ufunc_from_impl
+from ._decorators import deco_binary_ufunc_from_impl
 from ._detail import _ufunc_impl
 
 #
@@ -10,44 +10,44 @@
 #
 
 # 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

@@ -22,15 +22,11 @@ def wrapped(*args, dtype=None, **kwds):
 
 
 def emulate_out_arg(func):
-    """Simulate the out=... handling *for functions which do not need it*.
+    """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.
     """
     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)
 
@@ -45,10 +41,7 @@ def out_shape_dtype(func):
     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)
@@ -70,7 +63,7 @@ def wrapped(x1, *args, **kwds):
 
 # TODO: deduplicate with _ndarray/asarray_replacer,
 # and _wrapper/concatenate et al
-def deco_ufunc_from_impl(impl_func):
+def deco_binary_ufunc_from_impl(impl_func):
     @functools.wraps(impl_func)
     @dtype_to_torch
     @out_shape_dtype
@@ -98,7 +91,7 @@ def axis_keepdims_wrapper(func):
     # 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):
+    def wrapped(a, axis=None, keepdims=NoValue, *args, **kwds):
         from ._ndarray import ndarray, asarray
         tensor = asarray(a).get()
 

torch_np/_detail/_reductions.py

@@ -20,12 +20,9 @@ def _atleast_float(dtype, other_dtype):
     """
     if dtype is None:
         dtype = other_dtype
-    if not dtype.is_floating_point:
-        if dtype.is_complex:
-            pass  # pass through complex as is
-        else:
-            sctype = _scalar_types.default_float_type
-            dtype = sctype.torch_dtype
+    if not (dtype.is_floating_point or dtype.is_complex):
+        sctype = _scalar_types.default_float_type
+        dtype = sctype.torch_dtype
     return dtype
 
 
@@ -38,19 +35,19 @@ def count_nonzero(a, axis=None):
     return tensor
 
 
-def argmax(tensor, axis=None, out=None, *, keepdims=NoValue):
+def argmax(tensor, axis=None):
     axis = _util.allow_only_single_axis(axis)
     tensor = torch.argmax(tensor, axis)
     return tensor
 
-def argmin(tensor, axis=None, out=None, *, keepdims=NoValue):
+def argmin(tensor, axis=None):
     axis = _util.allow_only_single_axis(axis)
     tensor = torch.argmin(tensor, axis)
     return tensor
 
 
 def any(tensor, axis=None, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
 
     axis = _util.allow_only_single_axis(axis)
@@ -63,7 +60,7 @@ def any(tensor, axis=None, *, where=NoValue):
 
 
 def all(tensor, axis=None, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
 
     axis = _util.allow_only_single_axis(axis)
@@ -76,23 +73,23 @@ def all(tensor, axis=None, *, where=NoValue):
 
 
 def max(tensor, axis=None, initial=NoValue, where=NoValue):
-    if initial is not None or where is not None:
+    if initial is not NoValue or where is not NoValue:
         raise NotImplementedError
 
     result = tensor.amax(axis)
     return result
 
 
 def min(tensor, axis=None, initial=NoValue, where=NoValue):
-    if initial is not None or where is not None:
+    if initial is not NoValue or where is not NoValue:
         raise NotImplementedError
 
     result = tensor.amin(axis)
     return result
 
 
 def sum(tensor, axis=None, dtype=None, initial=NoValue, where=NoValue):
-    if initial is not None or where is not None:
+    if initial is not NoValue or where is not NoValue:
         raise NotImplementedError
 
     assert dtype is None or isinstance(dtype, torch.dtype)
@@ -109,7 +106,7 @@ def sum(tensor, axis=None, dtype=None, initial=NoValue, where=NoValue):
 
 
 def prod(tensor, axis=None, dtype=None, initial=NoValue, where=NoValue):
-    if initial is not None or where is not None:
+    if initial is not NoValue or where is not NoValue:
         raise NotImplementedError
 
     axis = _util.allow_only_single_axis(axis)
@@ -126,7 +123,7 @@ def prod(tensor, axis=None, dtype=None, initial=NoValue, where=NoValue):
 
 
 def mean(tensor, axis=None, dtype=None, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
 
     dtype = _atleast_float(dtype, tensor.dtype)
@@ -140,24 +137,26 @@ def mean(tensor, axis=None, dtype=None, *, where=NoValue):
 
 
 def std(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
 
     dtype = _atleast_float(dtype, tensor.dtype)
 
-    tensor = tensor.to(dtype)
+    if dtype is not None:
+        tensor = tensor.to(dtype)
     result = tensor.std(dim=axis, correction=ddof)
 
     return result
 
 
 def var(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
 
     dtype = _atleast_float(dtype, tensor.dtype)
 
-    tensor = tensor.to(dtype)
+    if dtype is not None:
+        tensor = tensor.to(dtype)
     result = tensor.var(dim=axis, correction=ddof)
 
     return result

torch_np/_detail/_util.py

@@ -54,8 +54,6 @@ def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
     By default, this forbids axes from being specified multiple times.
     Used internally by multi-axis-checking logic.
 
-    .. versionadded:: 1.13.0
-
     Parameters
     ----------
     axis : int, iterable of int
@@ -73,17 +71,6 @@ def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
     -------
     normalized_axes : tuple of int
         The normalized axis index, such that `0 <= normalized_axis < ndim`
-
-    Raises
-    ------
-    AxisError
-        If any axis provided is out of range
-    ValueError
-        If an axis is repeated
-
-    See also
-    --------
-    normalize_axis_index : normalizing a single scalar axis
     """
     # Optimization to speed-up the most common cases.
     if type(axis) not in (tuple, list):

torch_np/_helpers.py

@@ -40,10 +40,11 @@ def cast_and_broadcast(tensors, out, casting):
     return tuple(tensors)
 
 
-
 def result_or_out(result_tensor, out_array=None):
     """A helper for returns with out= argument."""
     if out_array is not None:
+        if not isinstance(out_array, ndarray):
+            raise TypeError("Return arrays must be of ArrayType")
         if result_tensor.shape != out_array.shape:
             raise ValueError("Bad size of the out array.")
         out_tensor = out_array.get()

torch_np/_wrapper.py

@@ -643,7 +643,7 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *, where=N
 
 @asarray_replacer()
 def nanmean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoValue):
-    if where is not None:
+    if where is not NoValue:
         raise NotImplementedError
     if dtype is None:
         dtype = a.dtype