Update lift to use the new PRNGKey check to suppress deprecation warn…

…ings PiperOrigin-RevId: 566339580
google · Sep 18, 2023 · 830d335 · 830d335
1 parent 5d846a5
commit 830d335
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/lift.py b/flax/core/lift.py
@@ -61,9 +61,9 @@
 
 
 def tree_map_rngs(fn, tree):
-  """Needed for mapping JAX random.* functions over KeyArray leaves."""
+  """Needed for mapping JAX random.* functions over PRNGKey leaves."""
   return jax.tree_util.tree_map(
-      fn, tree, is_leaf=lambda x: isinstance(x, random.KeyArray)
+      fn, tree, is_leaf=lambda x: jax.dtypes.issubdtype(x, jax.dtypes.prng_key)
   )