code: def get_dataset_tfds(
      dataset: str,
      split: str,
      shuffle_files: bool = True,
      data_dir: Optional[str] = None,
      skip_decode: Optional[Union[Sequence[str], Dict[Any, Any]]] = (
          "image",
      ),
  ):
      """Data provider."""
      builder = get_builder(dataset, data_dir)
      split = tfds.even_splits(
          split, jax.process_count(), drop_remainder=True
      )[jax.process_index()]
      skip_decoders = make_skip_decoders(
          skip_decode, builder.info.features
      )
      # Each host is responsible for a fixed subset of data
      return builder.as_dataset(
          split=split,
          shuffle_files=shuffle_files,
          read_config=tfds.ReadConfig(
              skip_prefetch=True,  # We prefetch after pipeline.
              try_autocache=False,  # We control this, esp. for few-shot.
              add_tfds_id=True,
          ),
          decoders=skip_decoders,
      )
  
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 16 for calling a global function with example usage: get_dataset_tfds(dataset, split, shuffle_files, data_dir, skip_decode) and returns: builder