From 9a1496fd61cdac50177c67c5090d7b542e41d347 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Fri, 8 Sep 2023 07:47:17 -0700
Subject: [PATCH] Update handling of typed PRNG keys

This follows the recommendations of [JEP 9263](https://github.com/google/jax/pull/17297)
---
 flax/core/scope.py | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/flax/core/scope.py b/flax/core/scope.py
index 2ad5370336..8848596dff 100644
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -1189,19 +1189,24 @@ def _is_valid_variables(variables: VariableDict) -> bool:
 
 def _is_valid_rng(rng: Array):
   """Checks whether rng is a valid JAX PRNGKey, also handling custom prngs."""
-  # New-style JAX KeyArrays have a base type.
-  if jax_config.jax_enable_custom_prng:  # type: ignore[attr-defined]
-    if not isinstance(rng, jax.random.KeyArray):
-      return False
-  # Old-style JAX PRNGKeys are plain uint32 arrays.
-  else:
-    if not isinstance(rng, (np.ndarray, jnp.ndarray)):
-      return False
-    if (
-        rng.shape != random.default_prng_impl().key_shape
-        or rng.dtype != jnp.uint32
-    ):
-      return False
+  # This check is valid for either new-style or old-style PRNG keys
+  if not isinstance(rng, (np.ndarray, jnp.ndarray)):
+    return False
+
+  # Handle new-style typed PRNG keys
+  if hasattr(jax.dtypes, 'prng_key'):  # JAX 0.4.14 or newer
+    if jax.dtypes.issubdtype(rng.dtype, jax.dtypes.prng_key):
+      return True
+  elif hasattr(jax.random, 'KeyArray'):  # Previous JAX versions
+    if isinstance(rng, jax.random.KeyArray):
+      return True
+
+  # Handle old-style raw PRNG keys
+  if (
+      rng.shape != random.default_prng_impl().key_shape
+      or rng.dtype != jnp.uint32
+  ):
+    return False
   return True