[flax:examples:mnist] Add get_fake_batch and get_apply_fn_and_args methods to train.py.

Preparing Flax examples for programmatic testing and benchmarking.

PiperOrigin-RevId: 815862166

Author: danielsuo

examples/mnist/main.py

@@ -22,12 +22,11 @@
 from absl import flags
 from absl import logging
 from clu import platform
+import train
 import jax
 from ml_collections import config_flags
 import tensorflow as tf
 
-import train
-
 
 FLAGS = flags.FLAGS
 

examples/mnist/train.py

@@ -20,6 +20,7 @@
 
 # See issue #620.
 # pytype: disable=wrong-keyword-args
+from typing import Any
 
 from absl import logging
 from flax import linen as nn
@@ -51,6 +52,38 @@ def __call__(self, x):
     return x
 
 
+def get_fake_batch(batch_size: int) -> Any:
+  """Returns fake data for the given batch size.
+
+  Args:
+    batch_size: The global batch size to generate.
+
+  Returns:
+    A properly sharded global batch of data.
+  """
+  rng = jax.random.PRNGKey(0)
+  images = jax.random.randint(rng, (batch_size, 28, 28, 1), 0, 255, jnp.uint8)
+  labels = jax.random.randint(rng, (batch_size,), 0, 10, jnp.int32)
+  return images, labels
+
+
+def get_apply_fn_and_args(
+    config: ml_collections.ConfigDict,
+) -> tuple[Any, tuple[Any, ...], dict[str, Any], tuple[Any, ...]]:
+  """Returns the apply function and args for the given config.
+
+  Args:
+    config: The training configuration.
+
+  Returns:
+    A tuple of the apply function, args and kwargs for the apply function, and
+    any metadata the training loop needs.
+  """
+  state = create_train_state(jax.random.key(0), config)
+  batch = get_fake_batch(config.batch_size)
+  return apply_model, (state, *batch), dict(), ()
+
+
 @jax.jit
 def apply_model(state, images, labels):
   """Computes gradients, loss and accuracy for a single batch."""
@@ -145,7 +178,7 @@ def train_and_evaluate(
         state, test_ds['image'], test_ds['label']
     )
 
-    logging.info(
+    logging.info(  # pytype: disable=logging-not-lazy
         'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, test_loss: %.4f,'
         ' test_accuracy: %.2f'
         % (