code: def load(
      split: Split,
      *,
      is_training: bool,
      batch_dims: Sequence[int],
      dtype: jnp.dtype = jnp.float32,
      transpose: bool = False,
      zeros: bool = False,
  ) -> Iterator[Batch]:
      """Loads the given split of the dataset."""
      if zeros:
          h, w, c = 224, 224, 3
          if transpose:
              image_dims = (*batch_dims[:-1], h, w, c, batch_dims[0])
          else:
              image_dims = (*batch_dims, h, w, c)
          batch = {
              "images": np.zeros(image_dims, dtype=dtype),
              "labels": np.zeros(batch_dims, dtype=np.uint32),
          }
          if is_training:
              yield from it.repeat(batch)
          else:
              num_batches = split.num_examples // np.prod(batch_dims)
              yield from it.repeat(batch, num_batches)
  
      if is_training:
          start, end = _shard(
              split, jax.process_index(), jax.process_count()
          )
      else:
          start, end = _shard(split, 0, 1)
      tfds_split = tfds.core.ReadInstruction(
          _to_tfds_split(split), from_=start, to=end, unit="abs"
      )
      ds = tfds.load(
          "imagenet2012:5.*.*",
          split=tfds_split,
          decoders={"image": tfds.decode.SkipDecoding()},
      )
  
      total_batch_size = np.prod(batch_dims)
  
      options = tf.data.Options()
      options.experimental_threading.private_threadpool_size = 48
      options.experimental_threading.max_intra_op_parallelism = 1
      options.experimental_optimization.map_parallelization = True
      if is_training:
          options.experimental_deterministic = False
      ds = ds.with_options(options)
  
      if is_training:
          if jax.process_count() > 1:
              # Only cache if we are reading a subset of the dataset.
              ds = ds.cache()
          ds = ds.repeat()
          ds = ds.shuffle(buffer_size=10 * total_batch_size, seed=0)
  
      else:
          if split.num_examples % total_batch_size != 0:
              raise ValueError(
                  f"Test/valid must be divisible by {total_batch_size}"
              )
  
      def preprocess(example):
          image = _preprocess_image(example["image"], is_training)
          label = tf.cast(example["label"], tf.int32)
          return {"images": image, "labels": label}
  
      ds = ds.map(
          preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
  
      def transpose_fn(batch):
          # We use the "double transpose trick" to improve performance for TPUs. Note
          # that this (typically) requires a matching HWCN->NHWC transpose in your
          # model code. The compiler cannot make this optimization for us since our
          # data pipeline and model are compiled separately.
          batch = dict(**batch)
          batch["images"] = tf.transpose(batch["images"], (1, 2, 3, 0))
          return batch
  
      def cast_fn(batch):
          batch = dict(**batch)
          batch["images"] = tf.cast(
              batch["images"], tf.dtypes.as_dtype(dtype)
          )
          return batch
  
      for i, batch_size in enumerate(reversed(batch_dims)):
          ds = ds.batch(batch_size)
          if i == 0:
              if transpose:
                  ds = ds.map(transpose_fn)  # NHWC -> HWCN
              # NOTE: You may be tempted to move the casting earlier on in the pipeline,
              # but for bf16 some operations will end up silently placed on the TPU and
              # this causes stalls while TF and JAX battle for the accelerator.
              if dtype != jnp.float32:
                  ds = ds.map(cast_fn)
  
      ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
      yield from tfds.as_numpy(ds)
  
imports: from collections.abc import Iterable, Iterator, Mapping, Sequence
  import enum
  import itertools
  import types
  from typing import Optional
  import jax
  import jax.numpy
  import numpy
  from packaging import version
  import tensorflow.compat.v2
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: load(split) and returns: dict, batch, batch