Repository: google/jax

Issue:
better errors for n-ary ops on incompatible types, especially key arrays
Operations on typed key arrays now err when we want them to, e.g. on `key + 7`. This guard's implementation relies on the failure case of general dtype promotion logic. The resulting error message can be confusing. Example:

```python
>>> jax.random.key(73) + 4
  ...
ValueError: dtype=key<fry> is not a valid dtype for JAX type promotion.
```

Here is another example, involving only keys, and an operation that's allowed on them when they match, but not on keys of different types:

```python
>>> jax.random.key(73) == jax.random.key(14)
Array(False, dtype=bool)
>>> jax.random.key(73) == jax.random.key(14, impl='rbg')
  ...
ValueError: dtype=key<rbg> is not a valid dtype for JAX type promotion.
```

It'd be nice to improve these error messages, maybe to something along the lines of "`<op>` is not defined on arguments of type `<A>` and `<B>`"

Would this also be an improvement to other cases where dtype promotion fails, more generally than with key arrays? Nice if so.


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -356,16 +356,33 @@ def check_no_float0s(fun_name: str, *args: Any):
 _check_no_float0s = check_no_float0s
 
 
+def check_for_prngkeys(fun_name: str, *args: Any):
+  """Check if args don't match and none of the args have typed prng dtype"""
+  arg_dtypes = [dtypes.dtype(arg) for arg in args]
+  if len(set(arg_dtypes)) < 2:
+    return  # Will be caught by extended dtype impl rules.
+  if any(dtypes.issubdtype(dt, dtypes.prng_key) for dt in arg_dtypes):
+    if len(arg_dtypes) == 1:
+      raise TypeError(
+        f"{fun_name} does not accept dtype {str(arg_dtypes[0])}.")
+    else:
+      raise TypeError(
+        f"{fun_name} does not accept dtypes {', '.join(map(str, arg_dtypes))}."
+      )
+
+
 def promote_args(fun_name: str, *args: ArrayLike) -> list[Array]:
   """Convenience function to apply Numpy argument shape and dtype promotion."""
   check_arraylike(fun_name, *args)
   _check_no_float0s(fun_name, *args)
+  check_for_prngkeys(fun_name, *args)
   return promote_shapes(fun_name, *promote_dtypes(*args))
 
 
 def promote_args_numeric(fun_name: str, *args: ArrayLike) -> list[Array]:
   check_arraylike(fun_name, *args)
   _check_no_float0s(fun_name, *args)
+  check_for_prngkeys(fun_name, *args)
   return promote_shapes(fun_name, *promote_dtypes_numeric(*args))
 
 
@@ -375,6 +392,7 @@ def promote_args_inexact(fun_name: str, *args: ArrayLike) -> list[Array]:
   Promotes non-inexact types to an inexact type."""
   check_arraylike(fun_name, *args)
   _check_no_float0s(fun_name, *args)
+  check_for_prngkeys(fun_name, *args)
   return promote_shapes(fun_name, *promote_dtypes_inexact(*args))
 
 


Test patch:
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -1250,7 +1250,7 @@ def test_vmap_split_mapped_key(self):
   def test_cannot_add(self):
     key = self.make_key(73)
     self.assertRaisesRegex(
-        ValueError, r'dtype=key<.*> is not a valid dtype for JAX type promotion.',
+        TypeError, r'add does not accept dtypes key<.*>, int.*',
         lambda: key + 47)
 
   def test_grad_of_prng_key(self):
@@ -1319,7 +1319,7 @@ def test_cannot_add(self):
     if not jnp.issubdtype(key.dtype, dtypes.prng_key):
       raise SkipTest('relies on typed key arrays')
     self.assertRaisesRegex(
-        ValueError, r'dtype=key<.*> is not a valid dtype for JAX type promotion.',
+        TypeError, r'add does not accept dtypes key<.*>, int.*',
         lambda: key + 47)
 
   def test_grad_of_prng_key(self):
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1348,9 +1348,9 @@ def test_scatter(self, idx):
 
   def test_errors(self):
     key = random.key(123)
-    with self.assertRaisesRegex(ValueError, "dtype=key<fry> is not a valid dtype"):
+    with self.assertRaisesRegex(TypeError, "add does not accept dtypes key<fry>, int.*"):
       jnp.add(key, 1)
-    with self.assertRaisesRegex(ValueError, "dtype=key<fry> is not a valid dtype"):
+    with self.assertRaisesRegex(TypeError, "add does not accept dtypes key<fry>, int.*"):
       key + 1
     with self.assertRaisesRegex(TypeError, "add does not accept dtype key<fry>"):
       jnp.add(key, key)