code: def train_epoch(
      train_step_fn: Callable[..., Tuple[TrainState, Metrics]],
      state: TrainState,
      train_batches: tf.data.Dataset,
      epoch: int,
      rngs: Optional[Dict[str, Any]] = None,
  ) -> Tuple[TrainState, Metrics]:
      """Train for a single epoch."""
      batch_metrics = []
      for batch in train_batches:
          batch = batch_to_numpy(batch)
          state, metrics = train_step_fn(state, batch, rngs)
          batch_metrics.append(metrics)
  
      # Compute the metrics for this epoch.
      batch_metrics = jax.device_get(batch_metrics)
      metrics = normalize_batch_metrics(batch_metrics)
  
      logging.info(
          "train epoch %03d loss %.4f accuracy %.2f",
          epoch,
          metrics.loss,
          metrics.accuracy * 100,
      )
  
      return state, metrics
  
imports: from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
  from absl import logging
  from flax import struct
  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
  import input_pipeline
  import models
label_desc: Use image node_id 11 for calling a global function with example usage: train_epoch(train_step_fn, state, train_batches, epoch, rngs) and returns: state,  metrics