code: def load(
      split: Split,
      is_training: bool,
      batch_dims: Sequence[int],
      tfds_data_dir: Optional[str] = None,
  ):
      """Loads the given split of the dataset."""
      if is_training:
          start, end = _shard(split, jax.host_id(), jax.host_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()},
          data_dir=tfds_data_dir,
      )
  
      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
      if is_training:
          options.experimental_deterministic = False
      ds = ds.with_options(options)
  
      if is_training:
          ds = ds.repeat()
          ds = ds.shuffle(buffer_size=10 * total_batch_size, seed=0)
  
      def preprocess(example):
          image = _preprocess_image(example["image"], is_training)
          image = tf.transpose(
              image, (2, 0, 1)
          )  # transpose HWC image to CHW format
          label = tf.cast(example["label"], tf.int32)
          return {"images": image, "labels": label}
  
      ds = ds.map(
          preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
  
      for batch_size in reversed(batch_dims):
          ds = ds.batch(batch_size)
  
      ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
  
      yield from tfds.as_numpy(ds)
  
file_path: objax/examples/image_classification/imagenet_resnet50_data.py
imports: import enum
  from typing import Optional, Sequence, Tuple
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: load(split, is_training, batch_dims, tfds_data_dir) and returns: dict