MAINT: stop attaching things to Tensors, move the logic to binary ufuncs

Author: ev-br

torch_np/_dtypes_impl.py

@@ -53,6 +53,18 @@ def result_type_impl(*tensors):
 
 # ### NEP 50 helpers ###
 
+SCALAR_TYPES = (int, bool, float, complex)
+
+
+def _dtype_for_scalar(py_type):
+    return {
+        bool: torch.bool,
+        int: torch.int64,
+        float: torch.float64,
+        complex: torch.complex128,
+    }[py_type]
+
+
 categories = [
     (torch.bool,),
     (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
@@ -69,3 +81,37 @@ def category(dtyp):
 
 
 dtype_for_cat = {0: torch.bool, 1: torch.int64, 2: torch.float64, 3: torch.complex128}
+
+
+def nep50_to_tensors(x1, x2):
+    x1_type, x2_type = type(x1), type(x2)
+    x1_is_weak = x1_type in SCALAR_TYPES
+    x2_is_weak = x2_type in SCALAR_TYPES
+    if x1_is_weak and x2_is_weak:
+        # two scalars: promote
+        x1 = torch.as_tensor(x1, dtype=_dtype_for_scalar(x1_type))
+        x2 = torch.as_tensor(x2, dtype=_dtype_for_scalar(x2_type))
+        return x1, x2
+    elif not (x1_is_weak or x2_is_weak):
+        # two tensors: nothing to do here
+        return x1, x2
+    else:
+        # scalar <op> scalar: NEP 50
+        weak, not_weak = (x1, x2) if x1_is_weak else (x2, x1)
+
+        # find the dtype for the weak's type
+        weak_dtype = _dtype_for_scalar(type(weak))
+
+        cat_weak = category(weak_dtype)
+        cat_not_weak = category(not_weak.dtype)
+
+        dt = not_weak.dtype if cat_weak <= cat_not_weak else dtype_for_cat[cat_weak]
+
+        # special-case complex + float32
+        if weak_dtype.is_complex and not_weak.dtype == torch.float32:
+            dt = torch.complex64
+
+        # finally, can cast make `weak` into a 0D tensor
+        weak = torch.as_tensor(weak, dtype=dt)
+
+        return (weak, not_weak) if x1_is_weak else (not_weak, weak)

torch_np/_ndarray.py

@@ -453,17 +453,7 @@ def __dlpack_device__(self):
 # The rest goes through asarray (preferred) or array.
 
 
-def _array(
-    obj,
-    dtype=None,
-    *,
-    copy=True,
-    order="K",
-    subok=False,
-    ndmin=0,
-    like=None,
-    is_weak=False,
-):
+def array(obj, dtype=None, *, copy=True, order="K", subok=False, ndmin=0, like=None):
     _util.subok_not_ok(like, subok)
     if order != "K":
         raise NotImplementedError
@@ -496,35 +486,12 @@ def _array(
     if dtype is not None:
         torch_dtype = _dtypes.dtype(dtype).torch_dtype
 
-    tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin, is_weak)
+    tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin)
     return ndarray(tensor)
 
 
-def array(obj, dtype=None, *, copy=True, order="K", subok=False, ndmin=0, like=None):
-    # The result of the public `np.array(obj)` is not weakly typed.
-    return _array(
-        obj,
-        dtype,
-        copy=copy,
-        order=order,
-        subok=subok,
-        ndmin=ndmin,
-        like=like,
-        is_weak=False,
-    )
-
-
-def _asarray(a, dtype=None, order="K", *, like=None, is_weak=False):
-    return _array(
-        a, dtype=dtype, order=order, like=like, copy=False, ndmin=0, is_weak=is_weak
-    )
-
-
 def asarray(a, dtype=None, order="K", *, like=None):
-    # The result of the public `np.asarray(obj)` is not weakly typed.
-    return _array(
-        a, dtype=dtype, order=order, like=like, copy=False, ndmin=0, is_weak=False
-    )
+    return array(a, dtype=dtype, order=order, like=like, copy=False, ndmin=0)
 
 
 def from_dlpack(x, /):

torch_np/_normalizations.py

@@ -9,9 +9,12 @@
 
 import torch
 
-from . import _dtypes, _util
+from . import _dtypes, _dtypes_impl, _util
 
 ArrayLike = typing.TypeVar("ArrayLike")
+Scalar = typing.Union[int, float, complex, bool]
+ArrayLikeOrScalar = typing.Union[ArrayLike, Scalar]
+
 DTypeLike = typing.TypeVar("DTypeLike")
 AxisLike = typing.TypeVar("AxisLike")
 NDArray = typing.TypeVar("NDarray")
@@ -38,15 +41,15 @@
 
 
 def normalize_array_like(x, parm=None):
-    from ._ndarray import _asarray
+    from ._ndarray import asarray
+
+    return asarray(x).tensor
 
-    # special case: python scalars are weakly typed
-    is_py_scalar = type(x) in (int, bool, float, complex)
-    if is_py_scalar:
-        dtype = _util._dtype_for_scalar(type(x))
-        return _asarray(x, dtype=dtype, is_weak=True).tensor
 
-    return _asarray(x).tensor
+def normalize_array_like_or_scalar(x, parm=None):
+    if type(x) in _dtypes_impl.SCALAR_TYPES:
+        return x
+    return normalize_array_like(x, parm)
 
 
 def normalize_optional_array_like(x, parm=None):
@@ -115,6 +118,7 @@ def normalize_casting(arg, parm=None):
 
 normalizers = {
     "ArrayLike": normalize_array_like,
+    "ArrayLike | Scalar": normalize_array_like_or_scalar,
     "Optional[ArrayLike]": normalize_optional_array_like,
     "Sequence[ArrayLike]": normalize_seq_array_like,
     "Optional[NDArray]": normalize_ndarray,

torch_np/_ufuncs.py

@@ -11,48 +11,11 @@
     DTypeLike,
     NotImplementedType,
     OutArray,
+    Scalar,
     normalizer,
 )
 
 
-def _ufunc_preprocess(
-    tensors, where, casting, order, dtype, subok, signature, extobj, scalars=False
-):
-
-    if scalars:
-        # if one of the original inputs is a weak scalar, activate the NEP 50 dance
-        # XXX: this is only needed for binops
-        x1, x2 = tensors
-        x1_is_weak = getattr(x1, "is_weakly_typed", False)
-        x2_is_weak = getattr(x2, "is_weakly_typed", False)
-        if x1_is_weak != x2_is_weak:
-            # scalar <op> array: NEP50; nothing to do otherwise
-            weak, non_weak = (x1, x2) if x1_is_weak else (x2, x1)
-
-            cat_weak = _dtypes_impl.category(weak.dtype)
-            cat_non_weak = _dtypes_impl.category(non_weak.dtype)
-
-            dt_weak = (
-                non_weak.dtype
-                if cat_weak <= cat_non_weak
-                else _dtypes_impl.dtype_for_cat[cat_weak]
-            )
-
-            # special-case complex + float32
-            if weak.dtype.is_complex and non_weak.dtype == torch.float32:
-                dt_weak = torch.complex64
-
-            weak = _util.cast_if_needed(weak, dt_weak)
-            tensors = (weak, non_weak) if x1_is_weak else (non_weak, weak)
-
-    if dtype is None:
-        dtype = _dtypes_impl.result_type_impl(*tensors)
-
-    tensors = _util.typecast_tensors(tensors, dtype, casting)
-
-    return tensors
-
-
 def _ufunc_postprocess(result, out, casting):
     if out is not None:
         result = _util.typecast_tensor(result, out.dtype.torch_dtype, casting)
@@ -78,8 +41,8 @@ def deco_binary_ufunc(torch_func):
 
     @normalizer
     def wrapped(
-        x1: ArrayLike,
-        x2: ArrayLike,
+        x1: ArrayLike | Scalar,
+        x2: ArrayLike | Scalar,
         /,
         out: Optional[OutArray] = None,
         *,
@@ -91,21 +54,16 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-        tensors = _ufunc_preprocess(
-            (x1, x2),
-            where,
-            casting,
-            order,
-            dtype,
-            subok,
-            signature,
-            extobj,
-            scalars=True,
-        )
-        result = torch_func(*tensors)
 
-        result = _ufunc_postprocess(result, out, casting)
+        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2)
+
+        if dtype is None:
+            dtype = _dtypes_impl.result_type_impl(x1, x2)
+        x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
 
+        result = torch_func(x1, x2)
+
+        result = _ufunc_postprocess(result, out, casting)
         return result
 
     wrapped.__qualname__ = torch_func.__name__
@@ -118,6 +76,7 @@ def wrapped(
 # matmul's signature is _slightly_ different from other ufuncs:
 # - no where=...
 # - additional axis=..., axes=...
+# - no NEP50 scalars in or out
 #
 @normalizer
 def matmul(
@@ -135,10 +94,12 @@ def matmul(
     axes: NotImplementedType = None,
     axis: NotImplementedType = None,
 ):
-    tensors = _ufunc_preprocess(
-        (x1, x2), True, casting, order, dtype, subok, signature, extobj
-    )
-    result = _binary_ufuncs_impl.matmul(*tensors)
+
+    if dtype is None:
+        dtype = _dtypes_impl.result_type_impl(x1, x2)
+    x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
+
+    result = _binary_ufuncs_impl.matmul(x1, x2)
 
     result = _ufunc_postprocess(result, out, casting)
     return result
@@ -178,11 +139,11 @@ def divmod(
     else:
         out1, out2 = out
 
-    tensors = _ufunc_preprocess(
-        (x1, x2), True, casting, order, dtype, subok, signature, extobj
-    )
+    if dtype is None:
+        dtype = _dtypes_impl.result_type_impl(x1, x2)
+    x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
 
-    quot, rem = _binary_ufuncs_impl.divmod(*tensors)
+    quot, rem = _binary_ufuncs_impl.divmod(x1, x2)
 
     quot = _ufunc_postprocess(quot, out1, casting)
     rem = _ufunc_postprocess(rem, out2, casting)

torch_np/_util.py

@@ -171,11 +171,7 @@ def typecast_tensors(tensors, target_dtype, casting):
     return tuple(typecast_tensor(t, target_dtype, casting) for t in tensors)
 
 
-def _dtype_for_scalar(py_type):
-    return {bool: torch.bool, int: torch.int64, float: torch.float64, complex: torch.complex128}[py_type]
-
-
-def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0, is_weak=False):
+def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
     """The core logic of the array(...) function.
 
     Parameters
@@ -205,22 +201,17 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0, is_weak=False):
     if isinstance(obj, torch.Tensor):
         tensor = obj
     else:
-        if is_weak:
-            # obj is a python scalar
-            dtype = dtype or _dtype_for_scalar(obj_type)
-            tensor = torch.as_tensor(obj, dtype=dtype)
-        else:
-            tensor = torch.as_tensor(obj)
+        tensor = torch.as_tensor(obj)
 
-            # tensor.dtype is the pytorch default, typically float32. If obj's elements
-            # are not exactly representable in float32, we've lost precision:
-            # >>> torch.as_tensor(1e12).item() - 1e12
-            # -4096.0
+        # tensor.dtype is the pytorch default, typically float32. If obj's elements
+        # are not exactly representable in float32, we've lost precision:
+        # >>> torch.as_tensor(1e12).item() - 1e12
+        # -4096.0
 
-            # Therefore, we treat `tensor.dtype` as a hint, and convert the
-            # original object *again*, this time with an explicit dtype.
-            torch_dtype = _dtypes_impl.get_default_dtype_for(tensor.dtype)
-            tensor = torch.as_tensor(obj, dtype=torch_dtype)
+        # Therefore, we treat `tensor.dtype` as a hint, and convert the
+        # original object *again*, this time with an explicit dtype.
+        torch_dtype = _dtypes_impl.get_default_dtype_for(tensor.dtype)
+        tensor = torch.as_tensor(obj, dtype=torch_dtype)
 
     # type cast if requested
     tensor = cast_if_needed(tensor, dtype)
@@ -234,6 +225,4 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0, is_weak=False):
     if copy:
         tensor = tensor.clone()
 
-    # Attach the flag *to the tensor* (will be used after normalizations)
-    tensor.is_weakly_typed = is_weak
     return tensor

torch_np/tests/test_nep50_examples.py

@@ -91,17 +91,3 @@ def test_nep50_exceptions(example):
 
         assert_allclose(result, new, atol=1e-16)
         assert result.dtype == new.dtype
-
-
-class TestScalarsWeakTyping:
-    def test_asarray_scalars(self):
-        assert tnp.asarray(3).tensor.is_weakly_typed is False
-
-    def test_asarray_asarray_scalars(self):
-        a = tnp.asarray(3)
-        assert tnp.asarray(a).tensor.is_weakly_typed is False
-
-    def test_scalar_scalar(self):
-        a = tnp.uint8(3)
-        is_weakly_typed = getattr(a.tensor, "is_weakly_typed", False)
-        assert is_weakly_typed is False