code: def load_split(
      ds_factory,
      batch_size: int,
      subset: Text = "train",
      num_frames: int = 32,
      stride: int = 2,
      num_test_clips: int = 1,
      min_resize: int = 256,
      crop_size: int = 224,
      one_hot_label: bool = True,
      zero_centering: bool = True,
      get_label_str: bool = False,
      augmentation_params: Optional[ml_collections.ConfigDict] = None,
      keep_key: bool = False,
      do_three_spatial_crops: bool = False,
      label_offset: int = 0,
  ) -> Tuple[tf.data.Dataset, int]:
      """Loads dataset using DMVR for pre-processing.
  
      DMVR dataset loader already does basic augmentation (random crop and flip in
        train mode. It also already shuffles and batches the data.
  
      Args:
        ds_factory: A DMVR factory to instantiate with the subset.
        batch_size: The batch_size to use.
        subset: train, validation or test
        num_frames: Number of frames per subclip.
        stride: Temporal stride to sample frames.
        num_test_clips: Number of test clips (1 by default). If more than 1, this
          will sample multiple linearly spaced clips within each video at test time.
          If 1, then a single clip in the middle of the video is sampled. The clips
          are aggreagated in the batch dimension.
        min_resize: Frames are resized so that min(height, width) is min_resize.
        crop_size: Final size of the frame after cropping the resized frames. Both
          height and width are the same.
        one_hot_label: If True, return one-hot version of the labels (ie [N, C])
          array. Otherwise, return [N]-array of labels.
        zero_centering: If True, frames are normalized to values in the interval
          [-1, 1]. If False, values are in the interval [0, 1].
        get_label_str: whether or not to return label as text.
          Note that strings cannot be used in pmapped functions in Jax!
        augmentation_params: Augmentation configurations in train mode.
        keep_key: bool; If true, also return the key for each example.
        do_three_spatial_crops: If true, take three spatial crops of each clip
          during testing.
        label_offset: If non-zero, this value is subtracted from the parsed label.
          Useful when dataset is 1-indexed.
  
      Returns:
        A pair `(ds, num_examples)` with
          ds: A `tf.data.Dataset` object
          num_examples: Number of examples in the dataset.
      """
      dataset = ds_factory(subset=subset).configure(
          is_training=(subset == "train"),
          num_frames=num_frames,
          stride=stride,
          num_test_clips=num_test_clips,
          min_resize=min_resize,
          crop_size=crop_size,
          zero_centering_image=zero_centering,
          one_hot_label=one_hot_label,
          get_label_str=get_label_str,
          label_offset=label_offset,
      )
  
      if subset == "train" and augmentation_params:
          dataset = video_ops.additional_augmentations(
              dataset,
              augmentation_params,
              crop_size,
              num_frames,
              zero_centering,
          )
  
      if subset != "train" and do_three_spatial_crops:
          rgb_feature_name = builders.IMAGE_FEATURE_NAME
  
          dataset.preprocessor_builder.replace_fn(
              f"{rgb_feature_name}_central_crop",
              functools.partial(
                  video_ops.three_spatial_crops, crop_size=crop_size
              ),
          )
  
          if num_test_clips == 1:
              # This means that reshaping is not part of the post-processing graph
              output_shape = (-1, num_frames, crop_size, crop_size, 3)
              dataset.postprocessor_builder.add_fn(
                  fn=lambda x: tf.reshape(x, output_shape),
                  feature_name=rgb_feature_name,
                  fn_name=f"{rgb_feature_name}_reshape",
              )
  
      logging.info(
          "Preprocessing graph: %s",
          dataset.preprocessor_builder.get_summary(),
      )
      logging.info(
          "Postprocessing graph: %s",
          dataset.postprocessor_builder.get_summary(),
      )
      num_examples = dataset.num_examples
  
      ds = dataset.make_dataset(
          batch_size=batch_size,
          shuffle=(subset == "train"),
          drop_remainder=(subset == "train"),
          keep_key=(subset != "train" and keep_key),
      )
  
      options = tf.data.Options()
      options.threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      return ds, num_examples
  
imports: import functools
  import os
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import builders
  from dmvr import modalities
  from dmvr import video_dataset
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.dataset_lib import video_ops
  from scenic.projects.vivit.data import file_utils
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: load_split(ds_factory, batch_size, subset, num_frames, stride, num_test_clips, min_resize, crop_size, one_hot_label, zero_centering, get_label_str, augmentation_params, keep_key, do_three_spatial_crops, label_offset) and returns: ds,  num_examples