From aa30300c6d987672886d0cd5581be319be265539 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 13 May 2023 19:53:49 +0300
Subject: [PATCH 01/17] ENH: introduce NEP 50 "weak scalars"

Make python scalars "weak": in type promotion, they do not type promote arrays:

   (np.int8(3) + 4).dtype == int8

- Note that array scalars (np.int8(3) etc) are 0D arrays, so they are not weak.
- Converting a weak scalar to an array (asarray(3) etc) makes it not weak.
- Scalars are only weak in ufuncs. In places like `np.dot([1, 2, 3], 4.0)`,
  the result is float64.
---
 torch_np/_ndarray.py                  | 39 ++++++++++++++++++++--
 torch_np/_normalizations.py           | 10 ++++--
 torch_np/_ufuncs.py                   |  8 +++++
 torch_np/_util.py                     | 38 ++++++++++++++--------
 torch_np/tests/test_nep50_examples.py | 47 ++++++++++++++++++++++-----
 5 files changed, 115 insertions(+), 27 deletions(-)

diff --git a/torch_np/_ndarray.py b/torch_np/_ndarray.py
index 41eac1f4..8ff6aad4 100644
--- a/torch_np/_ndarray.py
+++ b/torch_np/_ndarray.py
@@ -453,7 +453,17 @@ 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):
+def _array(
+    obj,
+    dtype=None,
+    *,
+    copy=True,
+    order="K",
+    subok=False,
+    ndmin=0,
+    like=None,
+    is_weak=False,
+):
     _util.subok_not_ok(like, subok)
     if order != "K":
         raise NotImplementedError
@@ -486,12 +496,35 @@ def array(obj, dtype=None, *, copy=True, order="K", subok=False, ndmin=0, like=N
     if dtype is not None:
         torch_dtype = _dtypes.dtype(dtype).torch_dtype
 
-    tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin)
+    tensor = _util._coerce_to_tensor(obj, torch_dtype, copy, ndmin, is_weak)
     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):
-    return array(a, dtype=dtype, order=order, like=like, copy=False, ndmin=0)
+    # 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
+    )
 
 
 def from_dlpack(x, /):
diff --git a/torch_np/_normalizations.py b/torch_np/_normalizations.py
index 9ab8932b..bb77200a 100644
--- a/torch_np/_normalizations.py
+++ b/torch_np/_normalizations.py
@@ -38,9 +38,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)
+    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_optional_array_like(x, parm=None):
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 8021d64a..66d8a6fa 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -68,6 +68,14 @@ def wrapped(
         result = torch_func(*tensors)
 
         result = _ufunc_postprocess(result, out, casting)
+
+        # if any of original inputs is weak, undo its effect on the result dtype
+        non_weaks = tuple(
+            x for x in (x1, x2) if not getattr(x, "is_weakly_typed", False)
+        )
+        if len(non_weaks) == 1:
+            result = _util.cast_if_needed(result, non_weaks[0].dtype)
+
         return result
 
     wrapped.__qualname__ = torch_func.__name__
diff --git a/torch_np/_util.py b/torch_np/_util.py
index fc7651dd..589fd2b5 100644
--- a/torch_np/_util.py
+++ b/torch_np/_util.py
@@ -171,7 +171,11 @@ def typecast_tensors(tensors, target_dtype, casting):
     return tuple(typecast_tensor(t, target_dtype, casting) for t in tensors)
 
 
-def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
+def _dtype_for_scalar(py_type):
+    return {bool: torch.bool, int: torch.int64, float: torch.float64}[py_type]
+
+
+def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0, is_weak=False):
     """The core logic of the array(...) function.
 
     Parameters
@@ -182,6 +186,10 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
         Coerce to this torch dtype
     copy : bool
         Copy or not
+    ndmin : int
+        The results as least this many dimensions
+    is_weak : bool
+        Whether obj is a weakly typed python scalar.
 
     Returns
     -------
@@ -197,20 +205,22 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
     if isinstance(obj, torch.Tensor):
         tensor = obj
     else:
-        tensor = torch.as_tensor(obj)
-        base = None
-
-        # At this point, `tensor.dtype` is the pytorch default. Our default may
-        # differ, so need to typecast. However, we cannot just do `tensor.to`,
-        # because if our desired dtype is wider then pytorch's, `tensor`
-        # may have lost precision:
+        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)
 
-        # int(torch.as_tensor(1e12)) - 1e12 equals -4096 (try it!)
+            # 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)
@@ -224,4 +234,6 @@ def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
     if copy:
         tensor = tensor.clone()
 
+    # Attach the flag *to the tensor* (will be used after normalizations)
+    tensor.is_weakly_typed = is_weak
     return tensor
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 37eb2adc..ed7ebc92 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -1,7 +1,20 @@
 """Test examples for NEP 50."""
 
+import torch_np as tnp
 from torch_np import array, float32, float64, inf, int64, uint8
-from torch_np.testing import assert_allclose
+
+# from torch_np.testing import assert_allclose
+
+
+def assert_allclose(actual, desired, rtol=1e-07, atol=0):
+    import torch
+
+    from torch_np import asarray
+
+    actual = asarray(actual).tensor
+    desired = asarray(desired).tensor
+    return torch.testing.assert_close(actual, desired, rtol=rtol, atol=atol)
+
 
 uint16 = uint8  # can be anything here, see below
 
@@ -44,16 +57,16 @@
 
 
 fails = [
-    "uint8(1) + 2",
-    "array([1], uint8) + 1",
-    "array([1], uint8) + 200",
-    "array([1], uint8) + array(1, int64)",
-    "array([100], uint8) + 200",
+    #    "uint8(1) + 2",
+    # "array([1], uint8) + 1",
+    # "array([1], uint8) + 200",
+    # "array([1], uint8) + array(1, int64)",
+    # "array([100], uint8) + 200",
     "array([1], uint8) + 300",
     "uint8(1) + 300",
-    "uint8(100) + 200",
-    "float32(1) + 3e100",
-    "array([1.], float32) + 3",
+    # "uint8(100) + 200",
+    # "float32(1) + 3e100",
+    # "array([1.], float32) + 3",
 ]
 
 
@@ -65,6 +78,8 @@ def test_nep50_exceptions(example):
 
     old, new = examples[example]
 
+    ###  breakpoint()
+
     if new == Exception:
         with assert_raises(OverflowError):
             eval(example)
@@ -77,3 +92,17 @@ 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

From b6246a0e23617647073a2afe3d18f2da938a7ae1 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Tue, 16 May 2023 20:23:55 +0300
Subject: [PATCH 02/17] actually implement NEP 50 weak scalar promotion

---
 torch_np/_dtypes_impl.py | 21 +++++++++++++++++++++
 torch_np/_ufuncs.py      | 34 +++++++++++++++++++++++++---------
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index f21b83be..66db9bb7 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -49,3 +49,24 @@ def result_type_impl(*tensors):
         dtyp = _cd._result_type_dict[dtyp][curr.dtype]
 
     return dtyp
+
+
+# ### NEP 50 helpers ###
+
+categories = [(torch.bool,),
+              (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
+              (torch.float16, torch.float32, torch.float64),
+              (torch.complex64, torch.complex128)]
+
+
+def category(dtyp):
+    for j, cat in enumerate(categories):
+        if dtyp in cat:
+            return j
+    raise ValueError(f"unknown dtype {dtyp}")
+
+
+dtype_for_cat = {0: torch.bool,
+                 1: torch.int64,
+                 2: torch.float64,
+                 3: torch.complex128}
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 66d8a6fa..36a8c09d 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -15,7 +15,30 @@
 )
 
 
-def _ufunc_preprocess(tensors, where, casting, order, dtype, subok, signature, extobj):
+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])
+
+            # TODO: special-case complex + float32
+
+            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)
 
@@ -63,19 +86,12 @@ def wrapped(
         extobj=None,
     ):
         tensors = _ufunc_preprocess(
-            (x1, x2), where, casting, order, dtype, subok, signature, extobj
+            (x1, x2), where, casting, order, dtype, subok, signature, extobj, scalars=True
         )
         result = torch_func(*tensors)
 
         result = _ufunc_postprocess(result, out, casting)
 
-        # if any of original inputs is weak, undo its effect on the result dtype
-        non_weaks = tuple(
-            x for x in (x1, x2) if not getattr(x, "is_weakly_typed", False)
-        )
-        if len(non_weaks) == 1:
-            result = _util.cast_if_needed(result, non_weaks[0].dtype)
-
         return result
 
     wrapped.__qualname__ = torch_func.__name__

From b4e5f73be0539b69e6a373e649bd9ac3a27ec4f9 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Tue, 16 May 2023 20:34:46 +0300
Subject: [PATCH 03/17] xfail a test: torch.uint8(-1) wraps around to 255

---
 .../tests/numpy_tests/core/test_scalarmath.py | 29 ++++++++++---------
 torch_np/tests/test_nep50_examples.py         |  8 -----
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/torch_np/tests/numpy_tests/core/test_scalarmath.py b/torch_np/tests/numpy_tests/core/test_scalarmath.py
index 13b10405..7b46c0c9 100644
--- a/torch_np/tests/numpy_tests/core/test_scalarmath.py
+++ b/torch_np/tests/numpy_tests/core/test_scalarmath.py
@@ -481,6 +481,22 @@ def test_numpy_scalar_relational_operators(self):
                 assert_(not np.array(1, dtype=dt1)[()] < np.array(0, dtype=dt2)[()],
                         "type %s and %s failed" % (dt1, dt2))
 
+        #Signed integers and floats
+        for dt1 in 'bhl' + np.typecodes['Float']:
+            assert_(1 > np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
+            assert_(not 1 < np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
+            assert_(-1 == np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
+
+            for dt2 in 'bhl' + np.typecodes['Float']:
+                assert_(np.array(1, dtype=dt1)[()] > np.array(-1, dtype=dt2)[()],
+                        "type %s and %s failed" % (dt1, dt2))
+                assert_(not np.array(1, dtype=dt1)[()] < np.array(-1, dtype=dt2)[()],
+                        "type %s and %s failed" % (dt1, dt2))
+                assert_(np.array(-1, dtype=dt1)[()] == np.array(-1, dtype=dt2)[()],
+                        "type %s and %s failed" % (dt1, dt2))
+
+    @pytest.mark.xfail(reason="pytorch unsigned wraps around")
+    def test_numpy_scalar_relational_operators_2(self):
         #Unsigned integers
         for dt1 in 'B':
             assert_(-1 < np.array(1, dtype=dt1)[()], "type %s failed" % (dt1,))
@@ -496,19 +512,6 @@ def test_numpy_scalar_relational_operators(self):
                 assert_(np.array(1, dtype=dt1)[()] != np.array(-1, dtype=dt2)[()],
                         "type %s and %s failed" % (dt1, dt2))
 
-        #Signed integers and floats
-        for dt1 in 'bhl' + np.typecodes['Float']:
-            assert_(1 > np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
-            assert_(not 1 < np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
-            assert_(-1 == np.array(-1, dtype=dt1)[()], "type %s failed" % (dt1,))
-
-            for dt2 in 'bhl' + np.typecodes['Float']:
-                assert_(np.array(1, dtype=dt1)[()] > np.array(-1, dtype=dt2)[()],
-                        "type %s and %s failed" % (dt1, dt2))
-                assert_(not np.array(1, dtype=dt1)[()] < np.array(-1, dtype=dt2)[()],
-                        "type %s and %s failed" % (dt1, dt2))
-                assert_(np.array(-1, dtype=dt1)[()] == np.array(-1, dtype=dt2)[()],
-                        "type %s and %s failed" % (dt1, dt2))
 
     def test_scalar_comparison_to_none(self):
         # Scalars should just return False and not give a warnings.
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index ed7ebc92..4d009563 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -57,16 +57,8 @@ def assert_allclose(actual, desired, rtol=1e-07, atol=0):
 
 
 fails = [
-    #    "uint8(1) + 2",
-    # "array([1], uint8) + 1",
-    # "array([1], uint8) + 200",
-    # "array([1], uint8) + array(1, int64)",
-    # "array([100], uint8) + 200",
     "array([1], uint8) + 300",
     "uint8(1) + 300",
-    # "uint8(100) + 200",
-    # "float32(1) + 3e100",
-    # "array([1.], float32) + 3",
 ]
 
 

From de0a611332ec223e4087c0c067f34abdb6fcd7c1 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Tue, 16 May 2023 20:35:12 +0300
Subject: [PATCH 04/17] lint

---
 torch_np/_dtypes_impl.py | 15 +++++++--------
 torch_np/_ufuncs.py      | 22 +++++++++++++++++-----
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 66db9bb7..a0c55317 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -53,10 +53,12 @@ def result_type_impl(*tensors):
 
 # ### NEP 50 helpers ###
 
-categories = [(torch.bool,),
-              (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
-              (torch.float16, torch.float32, torch.float64),
-              (torch.complex64, torch.complex128)]
+categories = [
+    (torch.bool,),
+    (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
+    (torch.float16, torch.float32, torch.float64),
+    (torch.complex64, torch.complex128),
+]
 
 
 def category(dtyp):
@@ -66,7 +68,4 @@ def category(dtyp):
     raise ValueError(f"unknown dtype {dtyp}")
 
 
-dtype_for_cat = {0: torch.bool,
-                 1: torch.int64,
-                 2: torch.float64,
-                 3: torch.complex128}
+dtype_for_cat = {0: torch.bool, 1: torch.int64, 2: torch.float64, 3: torch.complex128}
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 36a8c09d..3c0c5a89 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -15,7 +15,9 @@
 )
 
 
-def _ufunc_preprocess(tensors, where, casting, order, dtype, subok, signature, extobj, scalars=False):
+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
@@ -30,9 +32,11 @@ def _ufunc_preprocess(tensors, where, casting, order, dtype, subok, signature, e
             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])
+            dt_weak = (
+                non_weak.dtype
+                if cat_weak <= cat_non_weak
+                else _dtypes_impl.dtype_for_cat[cat_weak]
+            )
 
             # TODO: special-case complex + float32
 
@@ -86,7 +90,15 @@ def wrapped(
         extobj=None,
     ):
         tensors = _ufunc_preprocess(
-            (x1, x2), where, casting, order, dtype, subok, signature, extobj, scalars=True
+            (x1, x2),
+            where,
+            casting,
+            order,
+            dtype,
+            subok,
+            signature,
+            extobj,
+            scalars=True,
         )
         result = torch_func(*tensors)
 

From 299c047166dd71b0dd731838b2c105466395394f Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Tue, 16 May 2023 21:07:04 +0300
Subject: [PATCH 05/17] MAINT: remove cruft

---
 torch_np/tests/test_nep50_examples.py | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 4d009563..51fb04ee 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -2,18 +2,7 @@
 
 import torch_np as tnp
 from torch_np import array, float32, float64, inf, int64, uint8
-
-# from torch_np.testing import assert_allclose
-
-
-def assert_allclose(actual, desired, rtol=1e-07, atol=0):
-    import torch
-
-    from torch_np import asarray
-
-    actual = asarray(actual).tensor
-    desired = asarray(desired).tensor
-    return torch.testing.assert_close(actual, desired, rtol=rtol, atol=atol)
+from torch_np.testing import assert_allclose
 
 
 uint16 = uint8  # can be anything here, see below
@@ -49,7 +38,7 @@ def assert_allclose(actual, desired, rtol=1e-07, atol=0):
     "uint8(1) + 300": (int64(301), Exception),
     "uint8(100) + 200": (int64(301), uint8(44)),  # and RuntimeWarning
     "float32(1) + 3e100": (float64(3e100), float32(inf)),  # and RuntimeWarning [T7]
-    #      "array([0.1], float32) == 0.1": (array([False]),      unchanged),         # XXX: a typo in NEP50?
+    "array([0.1], float32) == 0.1": (array([False]),      unchanged),         # XXX: a typo in NEP50?
     "array([0.1], float32) == float64(0.1)": (array([True]), array([False])),
     "array([1.], float32) + 3": (array([4.0], float32), unchanged),
     "array([1.], float32) + int64(3)": (array([4.0], float32), array([4.0], float64)),
@@ -59,6 +48,7 @@ def assert_allclose(actual, desired, rtol=1e-07, atol=0):
 fails = [
     "array([1], uint8) + 300",
     "uint8(1) + 300",
+    "array([0.1], float32) == 0.1",
 ]
 
 
@@ -70,8 +60,6 @@ def test_nep50_exceptions(example):
 
     old, new = examples[example]
 
-    ###  breakpoint()
-
     if new == Exception:
         with assert_raises(OverflowError):
             eval(example)

From 6bcd35326f34d450abbd0b6b66458115a2a45570 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Tue, 16 May 2023 21:23:29 +0300
Subject: [PATCH 06/17] TST: add tests from the NEP text

---
 torch_np/tests/test_nep50_examples.py | 28 +++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 51fb04ee..850b01ec 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -1,10 +1,20 @@
 """Test examples for NEP 50."""
 
 import torch_np as tnp
-from torch_np import array, float32, float64, inf, int64, uint8
+from torch_np import (
+    array,
+    bool_,
+    complex64,
+    complex128,
+    float32,
+    float64,
+    inf,
+    int16,
+    int64,
+    uint8,
+)
 from torch_np.testing import assert_allclose
 
-
 uint16 = uint8  # can be anything here, see below
 
 
@@ -38,17 +48,27 @@
     "uint8(1) + 300": (int64(301), Exception),
     "uint8(100) + 200": (int64(301), uint8(44)),  # and RuntimeWarning
     "float32(1) + 3e100": (float64(3e100), float32(inf)),  # and RuntimeWarning [T7]
-    "array([0.1], float32) == 0.1": (array([False]),      unchanged),         # XXX: a typo in NEP50?
+    "array([0.1], float32) == 0.1": (
+        array([False]),
+        unchanged,
+    ),  # XXX: a typo in NEP50?
     "array([0.1], float32) == float64(0.1)": (array([True]), array([False])),
     "array([1.], float32) + 3": (array([4.0], float32), unchanged),
     "array([1.], float32) + int64(3)": (array([4.0], float32), array([4.0], float64)),
+    # additional examples from the NEP text
+    "int16(2) + 2": (int64(4), int16(4)),
+    "int16(4) + 4j": (complex128(4 + 4j), unchanged),
+    "float32(5) + 5j": (complex128(5 + 5j), complex64(5 + 5j)),
+    "bool_(True) + 1": (int64(2), unchanged),
+    "True + uint8(2)": (uint8(3), unchanged),
 }
 
 
 fails = [
     "array([1], uint8) + 300",
     "uint8(1) + 300",
-    "array([0.1], float32) == 0.1",
+    "array([0.1], float32) == 0.1",  # TODO: fix the example
+    "float32(5) + 5j",  # TODO: implement
 ]
 
 

From 3656962ad3c40d61fa32f0ef51f8ae03268b1348 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Wed, 17 May 2023 14:41:53 +0300
Subject: [PATCH 07/17] BUG: add a special case to 3j + float32(2) -> complex64

---
 torch_np/_normalizations.py           | 2 +-
 torch_np/_ufuncs.py                   | 4 +++-
 torch_np/_util.py                     | 2 +-
 torch_np/tests/test_nep50_examples.py | 1 -
 4 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/torch_np/_normalizations.py b/torch_np/_normalizations.py
index bb77200a..57f8a1bc 100644
--- a/torch_np/_normalizations.py
+++ b/torch_np/_normalizations.py
@@ -41,7 +41,7 @@ def normalize_array_like(x, parm=None):
     from ._ndarray import _asarray
 
     # special case: python scalars are weakly typed
-    is_py_scalar = type(x) in (int, bool, float)
+    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
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 3c0c5a89..03bf605b 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -38,7 +38,9 @@ def _ufunc_preprocess(
                 else _dtypes_impl.dtype_for_cat[cat_weak]
             )
 
-            # TODO: special-case complex + float32
+            # 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)
diff --git a/torch_np/_util.py b/torch_np/_util.py
index 589fd2b5..6e98af3b 100644
--- a/torch_np/_util.py
+++ b/torch_np/_util.py
@@ -172,7 +172,7 @@ def typecast_tensors(tensors, target_dtype, casting):
 
 
 def _dtype_for_scalar(py_type):
-    return {bool: torch.bool, int: torch.int64, float: torch.float64}[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):
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 850b01ec..d417e48d 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -68,7 +68,6 @@
     "array([1], uint8) + 300",
     "uint8(1) + 300",
     "array([0.1], float32) == 0.1",  # TODO: fix the example
-    "float32(5) + 5j",  # TODO: implement
 ]
 
 

From 19e96c2023025b41cba2d528c86304b16b159175 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Wed, 17 May 2023 17:11:48 +0300
Subject: [PATCH 08/17] MAINT: stop attaching things to Tensors, move the logic
 to binary ufuncs

---
 torch_np/_dtypes_impl.py              | 46 +++++++++++++++
 torch_np/_ndarray.py                  | 39 +------------
 torch_np/_normalizations.py           | 20 ++++---
 torch_np/_ufuncs.py                   | 83 +++++++--------------------
 torch_np/_util.py                     | 31 ++++------
 torch_np/tests/test_nep50_examples.py | 14 -----
 6 files changed, 93 insertions(+), 140 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index a0c55317..3166f6a0 100644
--- a/torch_np/_dtypes_impl.py
+++ b/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)
diff --git a/torch_np/_ndarray.py b/torch_np/_ndarray.py
index 8ff6aad4..41eac1f4 100644
--- a/torch_np/_ndarray.py
+++ b/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, /):
diff --git a/torch_np/_normalizations.py b/torch_np/_normalizations.py
index 57f8a1bc..abe075e4 100644
--- a/torch_np/_normalizations.py
+++ b/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,
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 03bf605b..2f2d45f8 100644
--- a/torch_np/_ufuncs.py
+++ b/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)
diff --git a/torch_np/_util.py b/torch_np/_util.py
index 6e98af3b..d3154d55 100644
--- a/torch_np/_util.py
+++ b/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
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index d417e48d..1b557413 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/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

From 4429b2fc7fa09bdd2679ba7faeecdfdbd71615d0 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Wed, 17 May 2023 17:43:20 +0300
Subject: [PATCH 09/17] BUG: detect uint8 overflow

---
 torch_np/_dtypes_impl.py              | 15 +++++++++++++--
 torch_np/tests/test_nep50_examples.py |  2 --
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 3166f6a0..e6c0f264 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -84,6 +84,12 @@ def category(dtyp):
 
 
 def nep50_to_tensors(x1, x2):
+    """If either of inputs is a python scalar, type-promote with NEP 50.
+
+    NB: NEP 50 mandates RuntimeWarnings on some overflows. We do not emit them:
+    we either raise OverflowError or just do the computation.
+    """
+
     x1_type, x2_type = type(x1), type(x2)
     x1_is_weak = x1_type in SCALAR_TYPES
     x2_is_weak = x2_type in SCALAR_TYPES
@@ -112,6 +118,11 @@ def nep50_to_tensors(x1, x2):
             dt = torch.complex64
 
         # finally, can cast make `weak` into a 0D tensor
-        weak = torch.as_tensor(weak, dtype=dt)
+        weak_ = torch.as_tensor(weak, dtype=dt)
+
+        # detect uint overflow: in PyTorch, uint8(-1) wraps around to 255,
+        # while NEP50 mandates an exception.
+        if weak_.dtype == torch.uint8 and weak_.item() != weak:
+            raise OverflowError(f"Python integer {weak} out of bounds for {weak_.dtype}")
 
-        return (weak, not_weak) if x1_is_weak else (not_weak, weak)
+        return (weak_, not_weak) if x1_is_weak else (not_weak, weak_)
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 1b557413..1a72cf5b 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -65,8 +65,6 @@
 
 
 fails = [
-    "array([1], uint8) + 300",
-    "uint8(1) + 300",
     "array([0.1], float32) == 0.1",  # TODO: fix the example
 ]
 

From ad8752a759dc30fac4b2d4e8c1276ede3c59b27f Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Wed, 17 May 2023 17:46:38 +0300
Subject: [PATCH 10/17] MAINT: improve comments, xfail messages

---
 torch_np/_dtypes_impl.py                           | 6 ++++--
 torch_np/tests/numpy_tests/core/test_scalarmath.py | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index e6c0f264..9826776b 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -87,7 +87,7 @@ def nep50_to_tensors(x1, x2):
     """If either of inputs is a python scalar, type-promote with NEP 50.
 
     NB: NEP 50 mandates RuntimeWarnings on some overflows. We do not emit them:
-    we either raise OverflowError or just do the computation.
+    we either raise an OverflowError or silently do the computation.
     """
 
     x1_type, x2_type = type(x1), type(x2)
@@ -123,6 +123,8 @@ def nep50_to_tensors(x1, x2):
         # detect uint overflow: in PyTorch, uint8(-1) wraps around to 255,
         # while NEP50 mandates an exception.
         if weak_.dtype == torch.uint8 and weak_.item() != weak:
-            raise OverflowError(f"Python integer {weak} out of bounds for {weak_.dtype}")
+            raise OverflowError(
+                f"Python integer {weak} out of bounds for {weak_.dtype}"
+            )
 
         return (weak_, not_weak) if x1_is_weak else (not_weak, weak_)
diff --git a/torch_np/tests/numpy_tests/core/test_scalarmath.py b/torch_np/tests/numpy_tests/core/test_scalarmath.py
index 7b46c0c9..4c145e1b 100644
--- a/torch_np/tests/numpy_tests/core/test_scalarmath.py
+++ b/torch_np/tests/numpy_tests/core/test_scalarmath.py
@@ -495,7 +495,7 @@ def test_numpy_scalar_relational_operators(self):
                 assert_(np.array(-1, dtype=dt1)[()] == np.array(-1, dtype=dt2)[()],
                         "type %s and %s failed" % (dt1, dt2))
 
-    @pytest.mark.xfail(reason="pytorch unsigned wraps around")
+    @pytest.mark.xfail(reason="NEP50")
     def test_numpy_scalar_relational_operators_2(self):
         #Unsigned integers
         for dt1 in 'B':

From 863ef4c3fb853f56d2d8bd5a5c345d0ce1725426 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 11:23:23 +0300
Subject: [PATCH 11/17] ENH: only turn on NEP50 for a subset of ufuncs

---
 torch_np/_dtypes_impl.py              | 41 ++++++++++++++++-----------
 torch_np/_normalizations.py           |  2 +-
 torch_np/_ufuncs.py                   | 12 ++++----
 torch_np/tests/test_nep50_examples.py | 35 +++++++++++++++++++++++
 4 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 9826776b..2f53c722 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -83,13 +83,12 @@ def category(dtyp):
 dtype_for_cat = {0: torch.bool, 1: torch.int64, 2: torch.float64, 3: torch.complex128}
 
 
-def nep50_to_tensors(x1, x2):
+def nep50_to_tensors(x1, x2, handle_weaks):
     """If either of inputs is a python scalar, type-promote with NEP 50.
 
     NB: NEP 50 mandates RuntimeWarnings on some overflows. We do not emit them:
     we either raise an OverflowError or silently do the computation.
     """
-
     x1_type, x2_type = type(x1), type(x2)
     x1_is_weak = x1_type in SCALAR_TYPES
     x2_is_weak = x2_type in SCALAR_TYPES
@@ -105,26 +104,34 @@ def nep50_to_tensors(x1, x2):
         # 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))
+        if handle_weaks:
+            # 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)
 
-        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]
 
-        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
 
-        # special-case complex + float32
-        if weak_dtype.is_complex and not_weak.dtype == torch.float32:
-            dt = torch.complex64
+            # detect overflows: in PyTorch, uint8(-1) wraps around to 255,
+            # while NEP50 mandates an exception.
+            if cat_weak == 1 and cat_not_weak == 1:
+                # integers
+                iinfo = torch.iinfo(not_weak.dtype)
+                if weak < iinfo.min or weak > iinfo.max: 
+                    raise OverflowError(
+                        f"Python integer {weak} out of bounds for {not_weak.dtype}"
+                    )
+
+        else:
+            # no NEP50 weak handling, fall back to the usual logic
+            dt = _dtype_for_scalar(type(weak))
 
         # finally, can cast make `weak` into a 0D tensor
         weak_ = torch.as_tensor(weak, dtype=dt)
 
-        # detect uint overflow: in PyTorch, uint8(-1) wraps around to 255,
-        # while NEP50 mandates an exception.
-        if weak_.dtype == torch.uint8 and weak_.item() != weak:
-            raise OverflowError(
-                f"Python integer {weak} out of bounds for {weak_.dtype}"
-            )
-
         return (weak_, not_weak) if x1_is_weak else (not_weak, weak_)
diff --git a/torch_np/_normalizations.py b/torch_np/_normalizations.py
index abe075e4..a52176ad 100644
--- a/torch_np/_normalizations.py
+++ b/torch_np/_normalizations.py
@@ -118,7 +118,7 @@ def normalize_casting(arg, parm=None):
 
 normalizers = {
     "ArrayLike": normalize_array_like,
-    "ArrayLike | Scalar": normalize_array_like_or_scalar,
+    "Union[ArrayLike, Scalar]": normalize_array_like_or_scalar,
     "Optional[ArrayLike]": normalize_optional_array_like,
     "Sequence[ArrayLike]": normalize_seq_array_like,
     "Optional[NDArray]": normalize_ndarray,
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 2f2d45f8..60601b92 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
@@ -32,6 +32,9 @@ def _ufunc_postprocess(result, out, casting):
 ]
 
 
+NEP50_FUNCS = ("add", "sub", "mul", "floordiv", "mod", "and", "or", "xor", "lshift", "rshift")
+
+
 def deco_binary_ufunc(torch_func):
     """Common infra for binary ufuncs.
 
@@ -41,8 +44,8 @@ def deco_binary_ufunc(torch_func):
 
     @normalizer
     def wrapped(
-        x1: ArrayLike | Scalar,
-        x2: ArrayLike | Scalar,
+        x1: Union[ArrayLike, Scalar],
+        x2: Union[ArrayLike, Scalar],
         /,
         out: Optional[OutArray] = None,
         *,
@@ -54,8 +57,7 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-
-        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2)
+        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2, torch_func.__name__ in NEP50_FUNCS)
 
         if dtype is None:
             dtype = _dtypes_impl.result_type_impl(x1, x2)
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 1a72cf5b..660903ab 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -1,5 +1,13 @@
 """Test examples for NEP 50."""
 
+import itertools
+
+try:
+    import numpy as _np
+    HAVE_NUMPY = True
+except ImportError:
+    HAVE_NUMPY = False
+
 import torch_np as tnp
 from torch_np import (
     array,
@@ -89,3 +97,30 @@ def test_nep50_exceptions(example):
 
         assert_allclose(result, new, atol=1e-16)
         assert result.dtype == new.dtype
+
+
+
+# ### Directly compare to numpy ###
+
+weaks = [True, 1, 2.0, 3j]
+non_weaks = [tnp.asarray(True),
+             tnp.uint8(1), tnp.int8(1), tnp.int32(1), tnp.int64(1),
+             tnp.float32(1), tnp.float64(1),
+             tnp.complex64(1), tnp.complex128(1)]
+
+@pytest.mark.skipif(not HAVE_NUMPY, reason="NumPy not found")
+@pytest.mark.parametrize("scalar, array", itertools.product(weaks, non_weaks))
+def test_direct_compare(scalar, array):
+    # compare to NumPy w/ NEP 50.
+    try:
+        state = _np._get_promotion_state()
+        _np._set_promotion_state("weak")
+
+        result = (scalar + array).tensor.numpy()
+        result_numpy = scalar + array.tensor.numpy()
+        assert result.dtype == result_numpy.dtype
+        assert result == result_numpy
+
+
+    finally:
+        _np._set_promotion_state(state)

From 22130af110f68f329662c64bb0283117a2080ed4 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 11:57:32 +0300
Subject: [PATCH 12/17] BUG: only use NEP50 if the default dtype is numpy's

---
 torch_np/_dtypes_impl.py | 11 +++++++----
 torch_np/_ufuncs.py      |  7 +++++--
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 2f53c722..d2fc20f1 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -90,6 +90,10 @@ def nep50_to_tensors(x1, x2, handle_weaks):
     we either raise an OverflowError or silently do the computation.
     """
     x1_type, x2_type = type(x1), type(x2)
+    if x1_type == torch.Tensor and x2_type == torch.Tensor:
+        # two tensors: nothing to do here
+        return x1, 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:
@@ -97,9 +101,6 @@ def nep50_to_tensors(x1, x2, handle_weaks):
         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)
@@ -128,8 +129,10 @@ def nep50_to_tensors(x1, x2, handle_weaks):
                     )
 
         else:
-            # no NEP50 weak handling, fall back to the usual logic
+            # no NEP50 weak handling, fall back to the usual logic---which
+            # includes looking up the default dtypes being numpy's or torch's
             dt = _dtype_for_scalar(type(weak))
+            dt = get_default_dtype_for(dt)
 
         # finally, can cast make `weak` into a 0D tensor
         weak_ = torch.as_tensor(weak, dtype=dt)
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 60601b92..6198a738 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -32,7 +32,8 @@ def _ufunc_postprocess(result, out, casting):
 ]
 
 
-NEP50_FUNCS = ("add", "sub", "mul", "floordiv", "mod", "and", "or", "xor", "lshift", "rshift")
+NEP50_FUNCS = ("add", "subtract", "multiply", "floor_divide", "remainder", "bitwise_and", "bitwise_or", "bitwise_xor", "left_shift", "right_shift")
+
 
 
 def deco_binary_ufunc(torch_func):
@@ -57,7 +58,9 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2, torch_func.__name__ in NEP50_FUNCS)
+        flag = (torch_func.__name__ in NEP50_FUNCS and
+                _dtypes_impl.default_dtypes == _dtypes_impl.default_dtypes_numpy)
+        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2, flag)
 
         if dtype is None:
             dtype = _dtypes_impl.result_type_impl(x1, x2)

From b7f035fb85ab410a2de1c04bc55e16a993854804 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 11:58:21 +0300
Subject: [PATCH 13/17] prostrate myself at the face of black

---
 torch_np/_dtypes_impl.py              |  2 +-
 torch_np/_ufuncs.py                   | 20 ++++++++++++++++----
 torch_np/tests/test_nep50_examples.py | 19 +++++++++++++------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index d2fc20f1..9a8bec2d 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -123,7 +123,7 @@ def nep50_to_tensors(x1, x2, handle_weaks):
             if cat_weak == 1 and cat_not_weak == 1:
                 # integers
                 iinfo = torch.iinfo(not_weak.dtype)
-                if weak < iinfo.min or weak > iinfo.max: 
+                if weak < iinfo.min or weak > iinfo.max:
                     raise OverflowError(
                         f"Python integer {weak} out of bounds for {not_weak.dtype}"
                     )
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index 6198a738..e58ccf02 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -32,8 +32,18 @@ def _ufunc_postprocess(result, out, casting):
 ]
 
 
-NEP50_FUNCS = ("add", "subtract", "multiply", "floor_divide", "remainder", "bitwise_and", "bitwise_or", "bitwise_xor", "left_shift", "right_shift")
-
+NEP50_FUNCS = (
+    "add",
+    "subtract",
+    "multiply",
+    "floor_divide",
+    "remainder",
+    "bitwise_and",
+    "bitwise_or",
+    "bitwise_xor",
+    "left_shift",
+    "right_shift",
+)
 
 
 def deco_binary_ufunc(torch_func):
@@ -58,8 +68,10 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-        flag = (torch_func.__name__ in NEP50_FUNCS and
-                _dtypes_impl.default_dtypes == _dtypes_impl.default_dtypes_numpy)
+        flag = (
+            torch_func.__name__ in NEP50_FUNCS
+            and _dtypes_impl.default_dtypes == _dtypes_impl.default_dtypes_numpy
+        )
         x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2, flag)
 
         if dtype is None:
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 660903ab..da812d63 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -4,6 +4,7 @@
 
 try:
     import numpy as _np
+
     HAVE_NUMPY = True
 except ImportError:
     HAVE_NUMPY = False
@@ -99,14 +100,21 @@ def test_nep50_exceptions(example):
         assert result.dtype == new.dtype
 
 
-
 # ### Directly compare to numpy ###
 
 weaks = [True, 1, 2.0, 3j]
-non_weaks = [tnp.asarray(True),
-             tnp.uint8(1), tnp.int8(1), tnp.int32(1), tnp.int64(1),
-             tnp.float32(1), tnp.float64(1),
-             tnp.complex64(1), tnp.complex128(1)]
+non_weaks = [
+    tnp.asarray(True),
+    tnp.uint8(1),
+    tnp.int8(1),
+    tnp.int32(1),
+    tnp.int64(1),
+    tnp.float32(1),
+    tnp.float64(1),
+    tnp.complex64(1),
+    tnp.complex128(1),
+]
+
 
 @pytest.mark.skipif(not HAVE_NUMPY, reason="NumPy not found")
 @pytest.mark.parametrize("scalar, array", itertools.product(weaks, non_weaks))
@@ -121,6 +129,5 @@ def test_direct_compare(scalar, array):
         assert result.dtype == result_numpy.dtype
         assert result == result_numpy
 
-
     finally:
         _np._set_promotion_state(state)

From 61322fdf985e83b23096b6bdc5c538488f460ca9 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 12:00:23 +0300
Subject: [PATCH 14/17] TST: a small bump of the FFT tolerance

---
 torch_np/tests/numpy_tests/fft/test_pocketfft.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torch_np/tests/numpy_tests/fft/test_pocketfft.py b/torch_np/tests/numpy_tests/fft/test_pocketfft.py
index c3534eb2..6ca7f3ab 100644
--- a/torch_np/tests/numpy_tests/fft/test_pocketfft.py
+++ b/torch_np/tests/numpy_tests/fft/test_pocketfft.py
@@ -44,8 +44,8 @@ def test_fft(self):
 
         np.random.seed(1234)
         x = random(30) + 1j*random(30)
-        assert_allclose(fft1(x), np.fft.fft(x), atol=2e-5)
-        assert_allclose(fft1(x), np.fft.fft(x, norm="backward"), atol=2e-5)
+        assert_allclose(fft1(x), np.fft.fft(x), atol=3e-5)
+        assert_allclose(fft1(x), np.fft.fft(x, norm="backward"), atol=3e-5)
         assert_allclose(fft1(x) / np.sqrt(30),
                         np.fft.fft(x, norm="ortho"), atol=5e-6)
         assert_allclose(fft1(x) / 30.,

From 2a5ac159037f89dd4a83e601d95107e2e5bb170e Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 12:31:57 +0300
Subject: [PATCH 15/17] DOC: a bit more on RuntimeWarnings

---
 torch_np/_dtypes_impl.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 9a8bec2d..81387d48 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -84,11 +84,7 @@ def category(dtyp):
 
 
 def nep50_to_tensors(x1, x2, handle_weaks):
-    """If either of inputs is a python scalar, type-promote with NEP 50.
-
-    NB: NEP 50 mandates RuntimeWarnings on some overflows. We do not emit them:
-    we either raise an OverflowError or silently do the computation.
-    """
+    """If either of inputs is a python scalar, type-promote with NEP 50."""
     x1_type, x2_type = type(x1), type(x2)
     if x1_type == torch.Tensor and x2_type == torch.Tensor:
         # two tensors: nothing to do here
@@ -120,6 +116,11 @@ def nep50_to_tensors(x1, x2, handle_weaks):
 
             # detect overflows: in PyTorch, uint8(-1) wraps around to 255,
             # while NEP50 mandates an exception.
+            #
+            # Note that we only check if each element of the binop overflows,
+            # not the result. Consider, e.g. `uint8(100) + 200`. Operands are OK
+            # in uint8, but the result overflows and wrap around 255.
+            # Numpy emits a RuntimeWarning, PyTorch does not, and we do not either.
             if cat_weak == 1 and cat_not_weak == 1:
                 # integers
                 iinfo = torch.iinfo(not_weak.dtype)

From 8dbe00032a6c9e7593cd6da9e2f4873f8b0103ed Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Thu, 18 May 2023 18:34:34 +0300
Subject: [PATCH 16/17] MAINT: add more binary ufuncs to the NEP 50 list, add a
 more comprehensive test

---
 torch_np/_ufuncs.py                           | 20 ++++++-
 .../tests/numpy_tests/core/test_scalarmath.py |  1 -
 torch_np/tests/test_nep50_examples.py         | 52 +++++++++++++++++++
 3 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index e58ccf02..bd8a1458 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -37,12 +37,28 @@ def _ufunc_postprocess(result, out, casting):
     "subtract",
     "multiply",
     "floor_divide",
+    "true_divide",
+    "divide",
     "remainder",
     "bitwise_and",
     "bitwise_or",
     "bitwise_xor",
-    "left_shift",
-    "right_shift",
+    "bitwise_left_shift",
+    "bitwise_right_shift",
+    "hypot",
+    "arctan2",
+    "logaddexp",
+    "logaddexp2",
+    "heaviside",
+    "copysign",
+    "fmax",
+    "minimum",
+    "fmin",
+    "maximum",
+    "fmod",
+    "gcd",
+    "lcm",
+    "pow",
 )
 
 
diff --git a/torch_np/tests/numpy_tests/core/test_scalarmath.py b/torch_np/tests/numpy_tests/core/test_scalarmath.py
index 4c145e1b..7c06e1ca 100644
--- a/torch_np/tests/numpy_tests/core/test_scalarmath.py
+++ b/torch_np/tests/numpy_tests/core/test_scalarmath.py
@@ -495,7 +495,6 @@ def test_numpy_scalar_relational_operators(self):
                 assert_(np.array(-1, dtype=dt1)[()] == np.array(-1, dtype=dt2)[()],
                         "type %s and %s failed" % (dt1, dt2))
 
-    @pytest.mark.xfail(reason="NEP50")
     def test_numpy_scalar_relational_operators_2(self):
         #Unsigned integers
         for dt1 in 'B':
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index da812d63..1e683d7e 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -131,3 +131,55 @@ def test_direct_compare(scalar, array):
 
     finally:
         _np._set_promotion_state(state)
+
+
+# ufunc name: [array.dtype]
+corners = {
+    "true_divide": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "divide": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "arctan2": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "copysign": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "heaviside": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "ldexp": ["bool_", "uint8", "int8", "int16", "int32", "int64"],
+    "power": ["uint8"],
+    "nextafter": ["float32"],
+}
+
+
+@pytest.mark.skipif(not HAVE_NUMPY, reason="NumPy not found")
+@pytest.mark.parametrize("name", tnp._ufuncs._binary)
+@pytest.mark.parametrize("scalar, array", itertools.product(weaks, non_weaks))
+def test_compare_ufuncs(name, scalar, array):
+
+    if name in corners and (
+        array.dtype.name in corners[name]
+        or tnp.asarray(scalar).dtype.name in corners[name]
+    ):
+        return pytest.skip(f"{name}(..., dtype=array.dtype)")
+
+    try:
+        state = _np._get_promotion_state()
+        _np._set_promotion_state("weak")
+
+        if name in ["matmul", "modf", "divmod"]:
+            return
+        ufunc = getattr(tnp, name)
+        ufunc_numpy = getattr(_np, name)
+
+        try:
+            result = ufunc(scalar, array)
+        except RuntimeError:
+            # RuntimeError: "bitwise_xor_cpu" not implemented for 'ComplexDouble' etc
+            result = None
+
+        try:
+            result_numpy = ufunc_numpy(scalar, array.tensor.numpy())
+        except TypeError:
+            # TypeError: ufunc 'hypot' not supported for the input types
+            result_numpy = None
+
+        if result is not None and result_numpy is not None:
+            assert result.tensor.numpy().dtype == result_numpy.dtype
+
+    finally:
+        _np._set_promotion_state(state)

From 2d1636e52e5adc2bebcb40ad91f81fa211b6d306 Mon Sep 17 00:00:00 2001
From: Mario Lezcano Casado <3291265+lezcano@users.noreply.github.com>
Date: Fri, 19 May 2023 07:23:55 +0100
Subject: [PATCH 17/17] ENH: NEP 50 "weak scalars" with dtype and PyTorch
 defaults (#143)

---
 torch_np/_dtypes_impl.py              | 133 +++++++++++++-------------
 torch_np/_ufuncs.py                   |  26 +++--
 torch_np/tests/test_nep50_examples.py |  40 ++++++--
 torch_np/tests/test_ufuncs_basic.py   |   4 -
 4 files changed, 115 insertions(+), 88 deletions(-)

diff --git a/torch_np/_dtypes_impl.py b/torch_np/_dtypes_impl.py
index 81387d48..fb53f721 100644
--- a/torch_np/_dtypes_impl.py
+++ b/torch_np/_dtypes_impl.py
@@ -53,7 +53,7 @@ def result_type_impl(*tensors):
 
 # ### NEP 50 helpers ###
 
-SCALAR_TYPES = (int, bool, float, complex)
+SCALAR_TYPES = {int, bool, float, complex}
 
 
 def _dtype_for_scalar(py_type):
@@ -65,77 +65,74 @@ def _dtype_for_scalar(py_type):
     }[py_type]
 
 
-categories = [
-    (torch.bool,),
-    (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64),
-    (torch.float16, torch.float32, torch.float64),
-    (torch.complex64, torch.complex128),
-]
+def _category(dtype):
+    return {
+        torch.bool: 0,
+        # int
+        torch.uint8: 1,
+        torch.int8: 1,
+        torch.int16: 1,
+        torch.int32: 1,
+        torch.int64: 1,
+        # float
+        torch.float16: 2,
+        torch.float32: 2,
+        torch.float64: 2,
+        # complex
+        torch.complex64: 3,
+        torch.complex128: 3,
+    }[dtype]
 
 
-def category(dtyp):
-    for j, cat in enumerate(categories):
-        if dtyp in cat:
-            return j
-    raise ValueError(f"unknown dtype {dtyp}")
+def nep50_to_tensors(x1, x2, handle_weaks):
+    """If either of inputs is a python scalar, type-promote with NEP 50."""
 
+    def to_tensor(scalar, dtype=None):
+        if dtype is None:
+            dtype = _dtype_for_scalar(type(scalar))
+            dtype = get_default_dtype_for(dtype)
+        return torch.as_tensor(scalar, dtype=dtype)
+
+    x1_is_weak = not isinstance(x1, torch.Tensor)
+    x2_is_weak = not isinstance(x2, torch.Tensor)
+    if not handle_weaks or (x1_is_weak and x2_is_weak):
+        x1 = to_tensor(x1) if x1_is_weak else x1
+        x2 = to_tensor(x2) if x2_is_weak else x2
+        return x1, x2
 
-dtype_for_cat = {0: torch.bool, 1: torch.int64, 2: torch.float64, 3: torch.complex128}
+    # scalar <op> tensor: NEP 50
+    assert x1_is_weak != x2_is_weak
 
+    weak, not_weak = (x1, x2) if x1_is_weak else (x2, x1)
 
-def nep50_to_tensors(x1, x2, handle_weaks):
-    """If either of inputs is a python scalar, type-promote with NEP 50."""
-    x1_type, x2_type = type(x1), type(x2)
-    if x1_type == torch.Tensor and x2_type == torch.Tensor:
-        # two tensors: nothing to do here
-        return x1, x2
+    # find the dtype for the weak's type
+    weak_dtype = _dtype_for_scalar(type(weak))
 
-    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
-    else:
-        # scalar <op> scalar: NEP 50
-        weak, not_weak = (x1, x2) if x1_is_weak else (x2, x1)
-
-        if handle_weaks:
-            # 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
-
-            # detect overflows: in PyTorch, uint8(-1) wraps around to 255,
-            # while NEP50 mandates an exception.
-            #
-            # Note that we only check if each element of the binop overflows,
-            # not the result. Consider, e.g. `uint8(100) + 200`. Operands are OK
-            # in uint8, but the result overflows and wrap around 255.
-            # Numpy emits a RuntimeWarning, PyTorch does not, and we do not either.
-            if cat_weak == 1 and cat_not_weak == 1:
-                # integers
-                iinfo = torch.iinfo(not_weak.dtype)
-                if weak < iinfo.min or weak > iinfo.max:
-                    raise OverflowError(
-                        f"Python integer {weak} out of bounds for {not_weak.dtype}"
-                    )
-
-        else:
-            # no NEP50 weak handling, fall back to the usual logic---which
-            # includes looking up the default dtypes being numpy's or torch's
-            dt = _dtype_for_scalar(type(weak))
-            dt = get_default_dtype_for(dt)
-
-        # 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_)
+    cat_weak = _category(weak_dtype)
+    cat_not_weak = _category(not_weak.dtype)
+
+    dt = not_weak.dtype if cat_weak <= cat_not_weak else None
+
+    # special-case complex + float32
+    if weak_dtype.is_complex and not_weak.dtype == torch.float32:
+        dt = torch.complex64
+
+    # detect overflows: in PyTorch, uint8(-1) wraps around to 255,
+    # while NEP50 mandates an exception.
+    #
+    # Note that we only check if each element of the binop overflows,
+    # not the result. Consider, e.g. `uint8(100) + 200`. Operands are OK
+    # in uint8, but the result overflows and wrap around 255.
+    # Numpy emits a RuntimeWarning, PyTorch does not, and we do not either.
+    if cat_weak == 1 and cat_not_weak == 1:
+        # integers
+        iinfo = torch.iinfo(not_weak.dtype)
+        if not (iinfo.min <= weak <= iinfo.max):
+            raise OverflowError(
+                f"Python integer {weak} out of bounds for {not_weak.dtype}"
+            )
+
+    # finally, can make `weak` into a 0D tensor
+    weak = to_tensor(weak, dt)
+
+    return (weak, not_weak) if x1_is_weak else (not_weak, weak)
diff --git a/torch_np/_ufuncs.py b/torch_np/_ufuncs.py
index bd8a1458..e6c983a7 100644
--- a/torch_np/_ufuncs.py
+++ b/torch_np/_ufuncs.py
@@ -84,20 +84,28 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-        flag = (
-            torch_func.__name__ in NEP50_FUNCS
-            and _dtypes_impl.default_dtypes == _dtypes_impl.default_dtypes_numpy
-        )
-        x1, x2 = _dtypes_impl.nep50_to_tensors(x1, x2, flag)
 
-        if dtype is None:
+        if dtype is not None:
+
+            def cast(x, dtype):
+                if isinstance(x, torch.Tensor):
+                    return _util.typecast_tensors((x,), dtype, casting)[0]
+                else:
+                    return torch.as_tensor(x, dtype=dtype)
+
+            x1 = cast(x1, dtype)
+            x2 = cast(x2, dtype)
+        elif isinstance(x1, torch.Tensor) and isinstance(x2, torch.Tensor):
             dtype = _dtypes_impl.result_type_impl(x1, x2)
-        x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
+            x1, x2 = _util.typecast_tensors((x1, x2), dtype, casting)
+        else:
+            x1, x2 = _dtypes_impl.nep50_to_tensors(
+                x1, x2, torch_func.__name__ in NEP50_FUNCS
+            )
 
         result = torch_func(x1, x2)
 
-        result = _ufunc_postprocess(result, out, casting)
-        return result
+        return _ufunc_postprocess(result, out, casting)
 
     wrapped.__qualname__ = torch_func.__name__
     wrapped.__name__ = torch_func.__name__
diff --git a/torch_np/tests/test_nep50_examples.py b/torch_np/tests/test_nep50_examples.py
index 1e683d7e..b87f3d64 100644
--- a/torch_np/tests/test_nep50_examples.py
+++ b/torch_np/tests/test_nep50_examples.py
@@ -102,8 +102,8 @@ def test_nep50_exceptions(example):
 
 # ### Directly compare to numpy ###
 
-weaks = [True, 1, 2.0, 3j]
-non_weaks = [
+weaks = (True, 1, 2.0, 3j)
+non_weaks = (
     tnp.asarray(True),
     tnp.uint8(1),
     tnp.int8(1),
@@ -113,19 +113,45 @@ def test_nep50_exceptions(example):
     tnp.float64(1),
     tnp.complex64(1),
     tnp.complex128(1),
-]
+)
+if HAVE_NUMPY:
+    dtypes = (
+        None,
+        _np.bool_,
+        _np.uint8,
+        _np.int8,
+        _np.int32,
+        _np.int64,
+        _np.float32,
+        _np.float64,
+        _np.complex64,
+        _np.complex128,
+    )
+else:
+    dtypes = (None,)
 
 
 @pytest.mark.skipif(not HAVE_NUMPY, reason="NumPy not found")
-@pytest.mark.parametrize("scalar, array", itertools.product(weaks, non_weaks))
-def test_direct_compare(scalar, array):
+@pytest.mark.parametrize(
+    "scalar, array, dtype", itertools.product(weaks, non_weaks, dtypes)
+)
+def test_direct_compare(scalar, array, dtype):
     # compare to NumPy w/ NEP 50.
     try:
         state = _np._get_promotion_state()
         _np._set_promotion_state("weak")
 
-        result = (scalar + array).tensor.numpy()
-        result_numpy = scalar + array.tensor.numpy()
+        if dtype is not None:
+            kwargs = {"dtype": dtype}
+        try:
+            result_numpy = _np.add(scalar, array.tensor.numpy(), **kwargs)
+        except Exception:
+            return
+
+        kwargs = {}
+        if dtype is not None:
+            kwargs = {"dtype": getattr(tnp, dtype.__name__)}
+        result = tnp.add(scalar, array, **kwargs).tensor.numpy()
         assert result.dtype == result_numpy.dtype
         assert result == result_numpy
 
diff --git a/torch_np/tests/test_ufuncs_basic.py b/torch_np/tests/test_ufuncs_basic.py
index 9fcb0ae8..02c8b312 100644
--- a/torch_np/tests/test_ufuncs_basic.py
+++ b/torch_np/tests/test_ufuncs_basic.py
@@ -380,10 +380,6 @@ def test_binary_ufunc_dtype(self):
         assert r32.dtype == "float32"
         assert r32 == 1
 
-        # casting of floating inputs to booleans
-        with assert_raises(TypeError):
-            np.add(1.0, 1e-15, dtype=bool)
-
         # now force the cast
         rb = np.add(1.0, 1e-15, dtype=bool, casting="unsafe")
         assert rb.dtype == bool