change func name in case it can avoid trace leaking

Author: Charles Li

src/maxtext/layers/nnx_decoders.py

@@ -505,7 +505,7 @@ def _extract_matching_state(template, full):
         return {k: _extract_matching_state(v, full[k]) for k, v in template.items()}
       return full
 
-    def layer_fn(carry, scanned_vars):
+    def _layer_fn(carry, scanned_vars):
       current_params, current_state = scanned_vars
 
       if self.config.parameter_memory_host_offload:
@@ -521,9 +521,9 @@ def layer_fn(carry, scanned_vars):
       # ONLY return non-param state to prevent memory duplication of weights
       return new_carry, new_current_state
 
-    layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+    checkpointed_layer_fn = jax.checkpoint(_layer_fn, policy=policy, prevent_cse=prevent_cse)
 
-    final_carry, scanned_other = jax.lax.scan(layer_fn, x_in, (params, state))
+    final_carry, scanned_other = jax.lax.scan(checkpointed_layer_fn, x_in, (params, state))
 
     if scan_axis != 0:
       params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), params)