code: def load_mnist(split: tfds.Split, batch_size: int):
      """Loads either training or test MNIST data.
  
      Args:
        split: either tfds.Split.TRAIN or tfds.Split.TEST.
  
      Returns:
        an iterator with pairs (images, labels). The images have shape
        (B, 28, 28, 1) and the labels have shape (B, 10), where B is the batch_size.
      """
      if _MOCK_DATA.value:
          with tfds.testing.mock_data(num_examples=batch_size):
              try:
                  ds = tfds.load("mnist", split=split)
              except Exception as e:
                  m = re.search(
                      r"metadata files were not found in (.+/)mnist/",
                      str(e),
                  )
                  if m:
                      msg = (
                          "TFDS mock_data is missing the mnist metadata files. Run the "
                          "`saved_model_main.py` binary and see where TFDS downloads "
                          "the mnist data set (typically ~/tensorflow_datasets/mnist). "
                          f"Copy the `mnist` directory to {m.group(1)} and re-run the test"
                      )
                      raise ValueError(msg) from e
                  else:
                      raise e
      else:
          ds = tfds.load("mnist", split=split)
  
      def _prepare_example(x):
          image = tf.cast(x["image"], tf.float32) / 255.0
          label = tf.one_hot(x["label"], 10)
          return (image, label)
  
      ds = ds.map(_prepare_example)
      # drop_remainder=True is important for use with Keras
      ds = (
          ds.cache()
          .shuffle(1000)
          .batch(batch_size, drop_remainder=True)
      )
      return ds
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import logging
  import re
  import time
  from typing import Any, Callable, Optional
  from absl import flags
  import flax
  from flax import linen
  import jax
  import jax.numpy
  from matplotlib import pyplot
  import numpy
  import optax
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: load_mnist(split, batch_size) and returns: ds, image,  label