code: def load_split(
      ds_factory,
      batch_size: int,
      shuffle_buffer_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,
      resize_keep_aspect_ratio: bool = True,
      one_hot_label: bool = True,
      zero_centering: bool = True,
      random_flip: bool = True,
      normalization_mean: Union[tf.Tensor, float] = 0.0,
      normalization_std: Union[tf.Tensor, float] = 1.0,
      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,
      train_frame_sampling_mode: str = "random",
      use_crop_and_resize_video_mae: bool = False,
      object_configs: ml_collections.ConfigDict = ml_collections.ConfigDict(),
      num_channels: int = 3,
  ) -> Tuple[tf.data.Dataset, int]:
      """Additionally support loading object boxes."""
      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,
          resize_keep_aspect_ratio=resize_keep_aspect_ratio,
          zero_centering_image=zero_centering,
          random_flip=random_flip,
          train_frame_sampling_mode=train_frame_sampling_mode,
          one_hot_label=one_hot_label,
          get_label_str=get_label_str,
          label_offset=label_offset,
          use_crop_and_resize_video_mae=use_crop_and_resize_video_mae,
          normalization_mean=normalization_mean,
          normalization_std=normalization_std,
          object_configs=object_configs,
      )
      del augmentation_params
  
      if (
          subset != "train"
          and do_three_spatial_crops
          and resize_keep_aspect_ratio
      ):
          with_boxes = object_configs.get("with_boxes", False)
          if with_boxes:
              dataset.preprocessor_builder.replace_fn(
                  f"{builders.IMAGE_FEATURE_NAME}_central_crop",
                  functools.partial(
                      objects_dataset_utils.three_spatial_crops_with_state,
                      crop_size=crop_size,
                  ),
              )
              dataset.preprocessor_builder.replace_fn(
                  "bboxes_central_crop",
                  functools.partial(
                      objects_dataset_utils.three_spatial_transform_box,
                      crop_size=crop_size,
                  ),
              )
          else:
              dataset.preprocessor_builder.replace_fn(
                  f"{builders.IMAGE_FEATURE_NAME}_central_crop",
                  functools.partial(
                      objects_dataset_utils.three_spatial_crops,
                      crop_size=crop_size,
                  ),
              )
  
          if num_test_clips == 1:
              # This means that reshaping is not part of the post-processing graph.
              dataset.postprocessor_builder.add_fn(
                  fn=lambda x: tf.reshape(  # pylint: disable=g-long-lambda
                      x,
                      (
                          -1,
                          num_frames,
                          crop_size,
                          crop_size,
                          num_channels,
                      ),
                  ),
                  feature_name=builders.IMAGE_FEATURE_NAME,
                  fn_name=f"{builders.IMAGE_FEATURE_NAME}_reshape",
              )
              if object_configs.get("return_boxes", -1) > 0:
                  dataset.postprocessor_builder.add_fn(
                      fn=lambda x: tf.reshape(  # pylint: disable=g-long-lambda
                          x,
                          (
                              -1,
                              num_frames,
                              object_configs.get("return_boxes"),
                              4,
                          ),
                      ),
                      feature_name="bboxes",
                      fn_name="bboxes_reshape",
                  )
  
      logging.info(
          "Frame sampling graph: %s",
          dataset.sampler_builder.get_summary(),
      )
      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
  
      if subset == "train":
          dataset.tune(shuffle_buffer=shuffle_buffer_size)
  
      # Validation and test splits are a single epoch, so that the last batch
      # is padded with zeroes. This is then repeated.
      ds = dataset.make_dataset(
          batch_size=batch_size,
          shuffle=(subset == "train"),
          num_epochs=None if (subset == "train") else 1,
          drop_remainder=(subset == "train"),
          keep_key=(subset != "train" and keep_key),
      )
  
      if subset != "train":
          ds = ds.repeat(None)
  
      options = tf.data.Options()
      options.threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      return ds, num_examples
  
imports: import functools
  from typing import Dict, Iterator, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import builders
  from dmvr import modalities
  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
  import scenic.projects.objectvivit.dataset_utils
  from scenic.projects.vivit.data.video_tfrecord_dataset import TFRecordDatasetFactory
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: load_split(ds_factory, batch_size, shuffle_buffer_size, subset, num_frames, stride, num_test_clips, min_resize, crop_size, resize_keep_aspect_ratio, one_hot_label, zero_centering, random_flip, normalization_mean, normalization_std, get_label_str, augmentation_params, keep_key, do_three_spatial_crops, label_offset, train_frame_sampling_mode, use_crop_and_resize_video_mae, object_configs, num_channels) and returns: ds,  num_examples