code: def train_epoch(state, train_ds, batch_size, rng):
      """Train for a single epoch."""
      train_ds_size = len(train_ds["image"])
      steps_per_epoch = train_ds_size // batch_size
  
      perms = jax.random.permutation(rng, len(train_ds["image"]))
      perms = perms[
          : steps_per_epoch * batch_size
      ]  # skip incomplete batch
      perms = perms.reshape((steps_per_epoch, batch_size))
  
      epoch_loss = []
      epoch_accuracy = []
  
      for perm in perms:
          batch_images = train_ds["image"][perm, ...]
          batch_labels = train_ds["label"][perm, ...]
          grads, loss, accuracy = apply_model(
              state, batch_images, batch_labels
          )
          state = update_model(state, grads)
          epoch_loss.append(loss)
          epoch_accuracy.append(accuracy)
      train_loss = np.mean(epoch_loss)
      train_accuracy = np.mean(epoch_accuracy)
      return state, train_loss, train_accuracy
  
file_path: flax/examples/mnist/train.py
imports: from absl import logging
  from flax import linen
  from flax.metrics import tensorboard
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: train_epoch(state, train_ds, batch_size, rng) and returns: state,  train_loss,  train_accuracy