Update to support JAX custom PRNG

JAX is in the progress of introducing customizable PRNGs (more details at google/jax#9263). This change updates neural_tangents PRNG handling to be compatible with both states of `jax_enable_custom_prng`.

PiperOrigin-RevId: 532574327

README.md

@@ -251,7 +251,7 @@ The `neural_tangents` (`nt`) package contains the following modules and function
 We remark the following differences between our library and the JAX one.
 
 * All `nt.stax` layers are instantiated with a function call, i.e. `nt.stax.Relu()` vs `jax.example_libraries.stax.Relu`.
-* All layers with trainable parameters use the [_NTK parameterization_](https://arxiv.org/abs/1806.07572) by default. However, `Dense` and `Conv` layers also support the [_standard parameterization_](https://arxiv.org/abs/2001.07301) via a `parameterization` keyword argument.
+* All layers with trainable parameters use the [_NTK parameterization_](https://arxiv.org/1806.07572) by default. However, `Dense` and `Conv` layers also support the [_standard parameterization_](https://arxiv.org/2001.07301) via a `parameterization` keyword argument.
 * `nt.stax` and `jax.example_libraries.stax` may have different layers and options available (for example `nt.stax` layers support `CIRCULAR` padding, have `LayerNorm`, but no `BatchNorm`.).
 
 

neural_tangents/_src/stax/requirements.py

@@ -811,7 +811,7 @@ def _propagate_shape(init_fn: InitFn,
   def init_and_apply(rng, x):
     _, params = init_fn(rng, tree_map(lambda x: x.shape, x))
     return apply_fn(params, x, rng=rng, **kwargs)
-  akey = ShapedArray((2,), np.uint32)
+  akey = jax.eval_shape(jax.random.PRNGKey, 0)
   try:
     shaped = eval_shape(init_and_apply, akey, shaped)
   except NotImplementedError:

neural_tangents/_src/utils/utils.py

@@ -270,7 +270,13 @@ def x1_is_x2(x1: np.ndarray,
   if jax.default_backend() == 'tpu':
     eps = 1e-4
 
-  return np.all(np.abs(x1 - x2) < eps)
+  try:
+    diff = x1 - x2
+  except TypeError:
+    # inputs are e.g. custom PRNGKeys which don't define subtraction.
+    return np.all(x1 == x2)
+  else:
+    return np.all(np.abs(diff) < eps)
 
 
 def _get_ndim(x: Union[int, Sized, np.ndarray]) -> int: