Fix maybe_extract_keys logic and move lambda outside of util.train …

…to avoid recompiling lax.cond.

PiperOrigin-RevId: 568289058

coix/util.py

@@ -57,13 +57,16 @@ def is_observed_site(site):
 
 
 def can_extract_key(args):
-  return args and (
-      jax.dtypes.issubdtype(args[0].dtype, jax.dtypes.prng_key)
-      or (
-          isinstance(args[0], jnp.ndarray)
-          and (args[0].dtype == jnp.uint32)
-          and (jnp.ndim(args[0]) >= 1)
-          and (args[0].shape[-1] == 2)
+  return (
+      args
+      and isinstance(args[0], jnp.ndarray)
+      and (
+          jax.dtypes.issubdtype(args[0].dtype, jax.dtypes.prng_key)
+          or (
+              (args[0].dtype == jnp.uint32)
+              and (jnp.ndim(args[0]) >= 1)
+              and (args[0].shape[-1] == 2)
+          )
       )
   )
 
@@ -119,6 +122,11 @@ def __call__(self, *args, **kwargs):
     return self.module.apply(self.params, *args, **kwargs)
 
 
+def _skip_update(grad, opt_state, params):
+  del params
+  return jax.tree_util.tree_map(jnp.zeros_like, grad), opt_state
+
+
 def train(
     loss_fn,
     init_params,
@@ -150,7 +158,7 @@ def step_fn(params, opt_state, *args, **kwargs):
     updates, opt_state = jax.lax.cond(
         jnp.isfinite(jax.flatten_util.ravel_pytree(grads)[0]).all(),
         optimizer.update,
-        lambda g, o, p: (jax.tree_util.tree_map(jnp.zeros_like, g), o),
+        _skip_update,
         grads,
         opt_state,
         params,