Allow using LazyRngs for flax init/apply.

This is occasionally useful for refactoring while preserving rng derivations.

PiperOrigin-RevId: 783598096

Author: levskaya

flax/core/scope.py

@@ -1184,6 +1184,10 @@ 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."""
+  # Allow for user-provided LazyRng - useful for compatibility when refactoring.
+  if isinstance(rng, LazyRng):
+    return True
+
   # This check is valid for either new-style or old-style PRNG keys
   if not isinstance(rng, (np.ndarray, jnp.ndarray)):
     return False

tests/core/core_scope_test.py

@@ -197,6 +197,10 @@ def test_rng_check_w_old_and_new_keys(self):
     self.assertTrue(scope._is_valid_rng(raw_key))
     self.assertFalse(scope._is_valid_rng(random.split(raw_key)))
 
+  def test_rng_check_w_lazy_rng(self):
+    key = random.key(0)
+    self.assertTrue(scope._is_valid_rng(scope.LazyRng.create(key, 1)))
+
   def test_jax_leak_detector(self):
     with jax.check_tracer_leaks(True):