Repository: google/jax

Issue:
Behavior of `!=` operator is confusing when rhs is a list
Compare:
```python
for np in jnp, onp:
  print(np)
  print(np.array([1,2]) != [1,2])
  print(np.array([1,2]) != [3,2])
  print(np.array([1,2]) != np.array([1,2]))
  print(np.array([1,2]) != np.array([3,2]))
```
Ouptut:
```
<module 'jax.numpy' from '/usr/local/lib/python3.6/dist-packages/jax/numpy/__init__.py'>
True
True
[False False]
[ True False]
<module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>
[False False]
[ True False]
[False False]
[ True False]
```

The strange dynamically-shaped result results in some strange downstream exceptions when dropping in jax.numpy in place of numpy, such as:
```python
print(any(onp.array([1,2]) != [1,2]))
print(any(jnp.array([1,2]) != [1,2]))
```
Output:
```
False
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-18-e89704e72507> in <module>()
      1 print(any(onp.array([1,2]) != [1,2]))
----> 2 print(any(jnp.array([1,2]) != [1,2]))

TypeError: 'bool' object is not iterable
```
(Of course the real use-case is `some_fn_returns_array(x) != some_fn_returns_list(y)`)


Discussion:
I think this is an artifact of jax.numpy not promoting lists to arrays automatically like numpy does. We're okay with that behavior in general, but perhaps the comparison functions like `equal`, `not_equal`, `less`, etc. could promote here.
The other alternative would be to at least raise rather than quietly returning `False`.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4575,20 +4575,22 @@ def wrapped(*args, **kwargs):
 
 _scalar_types = (int, float, complex, np.generic)
 _accepted_binop_types = (int, float, complex, np.generic, np.ndarray, ndarray)
+_rejected_binop_types = (list, tuple, set, dict)
 
-def _defer_to_unrecognized_arg(binary_op):
+def _defer_to_unrecognized_arg(opchar, binary_op, swap=False):
   # Ensure that other array types have the chance to override arithmetic.
   def deferring_binary_op(self, other):
     if hasattr(other, '__jax_array__'):
       other = other.__jax_array__()
-    if not isinstance(other, _accepted_binop_types):
-      return NotImplemented
-    return binary_op(self, other)
+    args = (other, self) if swap else (self, other)
+    if isinstance(other, _accepted_binop_types):
+      return binary_op(*args)
+    if isinstance(other, _rejected_binop_types):
+      raise TypeError(f"unsupported operand type(s) for {opchar}: "
+                      f"{type(args[0]).__name__!r} and {type(args[1]).__name__!r}")
+    return NotImplemented
   return deferring_binary_op
 
-def _swap_args(f):
-  return lambda x, y: f(y, x)
-
 def _unimplemented_setitem(self, i, x):
   msg = ("'{}' object does not support item assignment. JAX arrays are "
          "immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` "
@@ -4615,44 +4617,44 @@ def _deepcopy(self, memo):
     "deepcopy": _deepcopy,
     "neg": negative,
     "pos": positive,
-    "eq": _defer_to_unrecognized_arg(equal),
-    "ne": _defer_to_unrecognized_arg(not_equal),
-    "lt": _defer_to_unrecognized_arg(less),
-    "le": _defer_to_unrecognized_arg(less_equal),
-    "gt": _defer_to_unrecognized_arg(greater),
-    "ge": _defer_to_unrecognized_arg(greater_equal),
+    "eq": _defer_to_unrecognized_arg("==", equal),
+    "ne": _defer_to_unrecognized_arg("!=", not_equal),
+    "lt": _defer_to_unrecognized_arg("<", less),
+    "le": _defer_to_unrecognized_arg("<=", less_equal),
+    "gt": _defer_to_unrecognized_arg(">", greater),
+    "ge": _defer_to_unrecognized_arg(">=", greater_equal),
     "abs": abs,
-    "add": _defer_to_unrecognized_arg(add),
-    "radd": _defer_to_unrecognized_arg(add),
-    "sub": _defer_to_unrecognized_arg(subtract),
-    "rsub": _defer_to_unrecognized_arg(_swap_args(subtract)),
-    "mul": _defer_to_unrecognized_arg(multiply),
-    "rmul": _defer_to_unrecognized_arg(multiply),
-    "div": _defer_to_unrecognized_arg(divide),
-    "rdiv": _defer_to_unrecognized_arg(_swap_args(divide)),
-    "truediv": _defer_to_unrecognized_arg(true_divide),
-    "rtruediv": _defer_to_unrecognized_arg(_swap_args(true_divide)),
-    "floordiv": _defer_to_unrecognized_arg(floor_divide),
-    "rfloordiv": _defer_to_unrecognized_arg(_swap_args(floor_divide)),
-    "divmod": _defer_to_unrecognized_arg(divmod),
-    "rdivmod": _defer_to_unrecognized_arg(_swap_args(divmod)),
-    "mod": _defer_to_unrecognized_arg(mod),
-    "rmod": _defer_to_unrecognized_arg(_swap_args(mod)),
-    "pow": _defer_to_unrecognized_arg(power),
-    "rpow": _defer_to_unrecognized_arg(_swap_args(power)),
-    "matmul": _defer_to_unrecognized_arg(matmul),
-    "rmatmul": _defer_to_unrecognized_arg(_swap_args(matmul)),
-    "and": _defer_to_unrecognized_arg(bitwise_and),
-    "rand": _defer_to_unrecognized_arg(bitwise_and),
-    "or": _defer_to_unrecognized_arg(bitwise_or),
-    "ror": _defer_to_unrecognized_arg(bitwise_or),
-    "xor": _defer_to_unrecognized_arg(bitwise_xor),
-    "rxor": _defer_to_unrecognized_arg(bitwise_xor),
+    "add": _defer_to_unrecognized_arg("+", add),
+    "radd": _defer_to_unrecognized_arg("+", add, swap=True),
+    "sub": _defer_to_unrecognized_arg("-", subtract),
+    "rsub": _defer_to_unrecognized_arg("-", subtract, swap=True),
+    "mul": _defer_to_unrecognized_arg("*", multiply),
+    "rmul": _defer_to_unrecognized_arg("*", multiply, swap=True),
+    "div": _defer_to_unrecognized_arg("/", divide),
+    "rdiv": _defer_to_unrecognized_arg("/", divide, swap=True),
+    "truediv": _defer_to_unrecognized_arg("/", true_divide),
+    "rtruediv": _defer_to_unrecognized_arg("/", true_divide, swap=True),
+    "floordiv": _defer_to_unrecognized_arg("//", floor_divide),
+    "rfloordiv": _defer_to_unrecognized_arg("//", floor_divide, swap=True),
+    "divmod": _defer_to_unrecognized_arg("divmod", divmod),
+    "rdivmod": _defer_to_unrecognized_arg("divmod", divmod, swap=True),
+    "mod": _defer_to_unrecognized_arg("%", mod),
+    "rmod": _defer_to_unrecognized_arg("%", mod, swap=True),
+    "pow": _defer_to_unrecognized_arg("**", power),
+    "rpow": _defer_to_unrecognized_arg("**", power, swap=True),
+    "matmul": _defer_to_unrecognized_arg("@", matmul),
+    "rmatmul": _defer_to_unrecognized_arg("@", matmul, swap=True),
+    "and": _defer_to_unrecognized_arg("&", bitwise_and),
+    "rand": _defer_to_unrecognized_arg("&", bitwise_and, swap=True),
+    "or": _defer_to_unrecognized_arg("|", bitwise_or),
+    "ror": _defer_to_unrecognized_arg("|", bitwise_or, swap=True),
+    "xor": _defer_to_unrecognized_arg("^", bitwise_xor),
+    "rxor": _defer_to_unrecognized_arg("^", bitwise_xor, swap=True),
     "invert": bitwise_not,
-    "lshift": _defer_to_unrecognized_arg(left_shift),
-    "rshift": _defer_to_unrecognized_arg(right_shift),
-    "rlshift": _defer_to_unrecognized_arg(_swap_args(left_shift)),
-    "rrshift": _defer_to_unrecognized_arg(_swap_args(right_shift)),
+    "lshift": _defer_to_unrecognized_arg("<<", left_shift),
+    "rshift": _defer_to_unrecognized_arg(">>", right_shift),
+    "rlshift": _defer_to_unrecognized_arg("<<", left_shift, swap=True),
+    "rrshift": _defer_to_unrecognized_arg(">>", right_shift, swap=True),
     "round": _operator_round,
 }
 


Test patch:
diff --git a/tests/host_callback_test.py b/tests/host_callback_test.py
--- a/tests/host_callback_test.py
+++ b/tests/host_callback_test.py
@@ -941,8 +941,8 @@ def func(x):
               identity=True
               transforms=()
             ] b
-            _:f32[] = mul c 2.00
-            d:f32[] = mul 1.00 2.00
+            _:f32[] = mul 2.00 c
+            d:f32[] = mul 2.00 1.00
             e:f32[] = outside_call[
               arg_treedef={treedef}
               callback=...
@@ -960,8 +960,8 @@ def func(x):
               callback=...
               identity=True
             ] b
-            _:f32[] = mul c 2.00
-            d:f32[] = mul 1.00 2.00
+            _:f32[] = mul 2.00 c
+            d:f32[] = mul 2.00 1.00
             e:f32[] = mul d 3.00
           in (e,) }}""", jaxpr)
     assertMultiLineStrippedEqual(self, "", testing_stream.output)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -641,6 +641,36 @@ def testRightOperatorOverload(self, name, rng_factory, shapes, dtypes,
     with jtu.strict_promotion_if_dtypes_match(dtypes):
       self._CompileAndCheck( fun, args_maker, atol=tol, rtol=tol)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": f"{rec.test_name}_{othertype}", "name": rec.name, "othertype": othertype}
+    for rec in JAX_OPERATOR_OVERLOADS if rec.nargs == 2
+    for othertype in [dict, list, tuple, set]))
+  def testOperatorOverloadErrors(self, name, othertype):
+    # Test that binary operators with builtin collections raise a TypeError
+    # and report the types in the correct order.
+    data = [(1, 2), (2, 3)]
+    arr = jnp.array(data)
+    other = othertype(data)
+
+    msg = f"unsupported operand type.* 'DeviceArray' and '{othertype.__name__}'"
+    with self.assertRaisesRegex(TypeError, msg):
+      getattr(arr, name)(other)
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": f"{rec.test_name}_{othertype}", "name": rec.name, "othertype": othertype}
+    for rec in JAX_RIGHT_OPERATOR_OVERLOADS if rec.nargs == 2
+    for othertype in [dict, list, tuple, set]))
+  def testRightOperatorOverloadErrors(self, name, othertype):
+    # Test that binary operators with builtin collections raise a TypeError
+    # and report the types in the correct order.
+    data = [(1, 2), (2, 3)]
+    arr = jnp.array(data)
+    other = othertype(data)
+
+    msg = f"unsupported operand type.* '{othertype.__name__}' and 'DeviceArray'"
+    with self.assertRaisesRegex(TypeError, msg):
+      getattr(arr, name)(other)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": rec.test_name + f"_{dtype}",
        "rng_factory": rec.rng_factory,