struct.dataclass instances with jax.Array attributes cannot be compared

I have encountered a problem using `flax.strcut.dataclass` in my code. You can see this colab [here ](https://colab.research.google.com/drive/19LYY9RCqy2nUXGfnD-nDgCUJ8AO6J7i7?usp=sharing). I can recreate the bug with the following simple code:

```
from flax import struct, linen as nn
import jax

@struct.dataclass
class RNGClass:
rng: "Any"

r1 = RNGClass(jax.random.PRNGKey(0))
r2 = RNGClass(jax.random.PRNGKey(0))
r1 == r2
```

I expect this to return true. However , the two instances of the `RNGClass` cannot be compared returning: `ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()`.

I have tracked down the problem and it is because in [here](https://github.com/google/jax/blob/55fbe1c7b5a682dfe986a653a904921fa663df67/jax/_src/numpy/array_methods.py#L243) in jax it uses `jp.equal` to check the equality of the two rngs terms which returns a `[True, True]`
array.