Merge pull request #4639 from google:flaxlib-types

PiperOrigin-RevId: 750628750

Author: Flax Authors

benchmarks/nnx_simple_training.py

@@ -116,14 +116,11 @@ def test_step_nnx(model: MLP, batch):
       loss = jnp.mean((y - y_pred) ** 2)
       return {'loss': loss}
 
-    cached_train_step_nnx = nnx.cached_partial(train_step_nnx, model, optimizer)
-    cached_test_step_nnx = nnx.cached_partial(test_step_nnx, model)
-
     for step, batch in enumerate(dataset(X, Y, batch_size)):
-      cached_train_step_nnx(batch)
+      train_step_nnx(model, optimizer, batch)
 
       if step % 1000 == 0:
-        logs = cached_test_step_nnx((X, Y))
+        logs = test_step_nnx(model, (X, Y))
 
       if step >= total_steps - 1:
         break