Repository: google/jax

Issue:
NamedShape cannot be compared with None
Occasionally, when running `jax.random.bernoulli(rng_key)` on CI, I'm seeing the error
```
test/test_distributions.py:657: in gen_values_outside_bounds
    sign = random.bernoulli(key1)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

key = array([2986633504,  617412143], dtype=uint32)
p = DeviceArray(0.5, dtype=float32), shape = None

    def bernoulli(key: jnp.ndarray,
                  p: RealArray = np.float32(0.5),
                  shape: Optional[Union[Sequence[int], NamedShape]] = None) -> jnp.ndarray:
      """Sample Bernoulli random values with given shape and mean.
    
      Args:
        key: a PRNGKey used as the random key.
        p: optional, a float or array of floats for the mean of the random
          variables. Must be broadcast-compatible with ``shape``. Default 0.5.
        shape: optional, a tuple of nonnegative integers representing the result
          shape. Must be broadcast-compatible with ``p.shape``. The default (None)
          produces a result shape equal to ``p.shape``.
    
      Returns:
        A random array with boolean dtype and shape given by ``shape`` if ``shape``
        is not None, or else ``p.shape``.
      """
      dtype = dtypes.canonicalize_dtype(lax.dtype(p))
      if shape is not None:
        shape = core.as_named_shape(shape)
      if not jnp.issubdtype(dtype, np.floating):
        msg = "bernoulli probability `p` must have a floating dtype, got {}."
        raise TypeError(msg.format(dtype))
      p = lax.convert_element_type(p, dtype)
>     return _bernoulli(key, p, shape)  # type: ignore
E     ValueError: static arguments should be comparable using __eq__.The following error was raised when comparing two objects of types <class 'jax.core.NamedShape'> and <class 'NoneType'>. The error was:
E     TypeError: NamedShape doesn't support comparisons with <class 'NoneType'>
```
Looking deeper into the JAX API, I think that this is triggered when using `cpp_jit`. But it seems that under `cpp_jit`, a static argument can accept both non-eq argument and None, as illustrated in the following code
```python
import jax

class NoEq:  # mimic jax/tests/api_test.py
  def __hash__(self):
    return 1
  def __eq__(self, other):
    raise NotImplementedError("Cannot compare!")

f = lambda x, y: x + 3
jf = jax.jit(f, static_argnums=(1,))
jf(1, NoEq())
jf(1, None)  # pass
jf(1, 1)  # fail 
```

That makes me confused. I'm not sure why such comparison happens with `random.bernoulli`. It happens [occasionally on CI](https://github.com/pyro-ppl/numpyro/runs/4059917016?check_suite_focus=true) and to resolve it, I can rerun all the CI jobs. I can't replicate the issue locally so my assumption is such comparison only happens in some specific system. I guess a solution is to allow comparing NamedShape and None but this seems to be an issue of `cpp_jit`, rather than NamedShape.


Discussion:

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -1643,7 +1643,7 @@ def __eq__(self, other):
       return (self.__positional, self.__named) == (other.__positional, other.__named)
     if isinstance(other, tuple):
       return not self.__named and self.__positional == other
-    raise TypeError(f"NamedShape doesn't support comparisons with {type(other)}")
+    return False
 
   def __hash__(self):
     named = frozenset(self.__named.items())


Test patch:
diff --git a/tests/core_test.py b/tests/core_test.py
--- a/tests/core_test.py
+++ b/tests/core_test.py
@@ -515,6 +515,13 @@ def test_typecompat_named_shape(self):
     aval3 = core.ShapedArray((2, 3), np.float32, False, {'i': 5})
     self.assertFalse(core.typecompat(aval1, aval3))
 
+  def test_named_shape_comparision(self):
+    self.assertTrue(core.NamedShape(2, 3) == (2, 3))
+    self.assertFalse(core.NamedShape(2, i=3) == (2,))
+    self.assertFalse(core.NamedShape(2, i=3) == (2, 3))
+    self.assertFalse(core.NamedShape(2, i=3) == None)
+    self.assertFalse(core.NamedShape() == [])
+
 
 class DynamicShapesTest(jtu.JaxTestCase):