code: def load_split_from_tfds(
      dataset_name,
      batch_size,
      split,
      data_dir=None,
      preprocess_example=None,
      augment_train_example=None,
      postprocess_batch=None,
      shuffle_buffer_size=None,
      shuffle_seed=0,
      cache=True,
      **kwargs,
  ):
      """Loads a split from a dataset using TensorFlow Datasets.
  
      Args:
        dataset_name: str; Name of the dataset to be used to load from tfds.
        batch_size: int; The batch size returned by the data pipeline.
        split: str; Name of  the split to be loaded.
        data_dir: str; Data directory.
        preprocess_example: function; A function that given an example, returns the
          preprocessed example. Note that the preprocessing is done BEFORE caching
          to re-use them.
        augment_train_example: A function that given a train example returns the
          augmented example. Note that this function is applied AFTER caching and
          repeat to get true randomness.
        postprocess_batch: function; A function that given a batch, returns the
          postprocessed batch.
        shuffle_buffer_size: int; Size of the tf.data.dataset shuffle buffer.
        shuffle_seed: int; Seed for shuffling the training data.
        cache: bool; Whether to cache the dataset in memory.
        **kwargs: Passed to tfds.builder().
  
      Returns:
        A `tf.data.Dataset`, and dataset information.
      """
      return load_split_from_tfds_builder(
          builder=tfds.builder(
              dataset_name, data_dir=data_dir, **kwargs
          ),
          batch_size=batch_size,
          split=split,
          preprocess_example=preprocess_example,
          augment_train_example=augment_train_example,
          postprocess_batch=postprocess_batch,
          shuffle_buffer_size=shuffle_buffer_size,
          shuffle_seed=shuffle_seed,
          cache=cache,
      )
  
imports: import collections
  import dataclasses
  import functools
  import itertools
  from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Union
  from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 9 for calling a global function with example usage: load_split_from_tfds(dataset_name, batch_size, split, data_dir, preprocess_example, augment_train_example, postprocess_batch, shuffle_buffer_size, shuffle_seed, cache) and returns: load_split_from_tfds_builder