code: def load_dataset(
      split: str,
      *,
      shuffle: bool,
      batch_size: int,
  ) -> Iterator[Batch]:
      """Loads the dataset as a generator of batches."""
      ds, ds_info = tfds.load(
          "mnist:3.*.*", split=split, with_info=True
      )
      ds.cache()
      if shuffle:
          ds = ds.shuffle(ds_info.splits[split].num_examples, seed=0)
      ds = ds.repeat()
      ds = ds.batch(batch_size)
      return iter(tfds.as_numpy(ds))
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 9 for calling a global function with example usage: load_dataset(split) and returns: iter