code: def coco_load_split_from_tfds(
      batch_size,
      *,
      train,
      preprocess_fn,
      decode_fn,
      cache=False,
      max_size=1333,
      max_boxes=100,
      shuffle_buffer_size=1000,
      shuffle_seed=0,
  ):
      """Loads a split from the COCO dataset using TensorFlow Datasets.
  
      Args:
        batch_size: int; The batch size returned by the data pipeline.
        train: bool; Whether to load the train or evaluation split.
        preprocess_fn: function; A function that given an example, train flag,
          and dtype returns the preprocessed the example. Note that the
          preprocessing is done BEFORE caching to re-use them.
        decode_fn: A function that given an example decodes the image, converts
          it to float32, mean-subtracts it, and pulls out the relevant parts from
          the tfds features.
        cache: bool; whether to use the ds.cache or nor.
        max_size: int; Maximum image size.
        max_boxes: int; Maximum number of boxes.
        shuffle_buffer_size: int; Size of the shuffle buffer.
        shuffle_seed: int; Seed for shuffling the training data.
  
      Returns:
        A `tf.data.Dataset`, and dataset info.
      """
      split = "train" if train else "validation"
      builder = tfds.builder("coco/2017")
  
      # Each host is responsible for a fixed subset of data.
      data_range = tfds.even_splits(split, jax.process_count())[
          jax.process_index()
      ]
      ds = builder.as_dataset(split=data_range, shuffle_files=False)
      options = tf.data.Options()
      options.threading.private_threadpool_size = 48
      ds = ds.with_options(options)
      ds = ds.map(
          decode_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
      if cache:
          ds = ds.cache()
  
      # TLDR: make sure max_boxes is set >=64.
      # NOTE: the number of boxes/labels always needs to be strictly larger than 63
      #   to ensure that there is at least one dummy target corresponding
      #   to an empty bounding box, and that the last target box is such a dummy
      #   empty target. This is needed for matching functions that in principle only
      #   produce matches with non-empty target boxes, and produce dummy matches
      #   with an empty target for the rest of the unmatched predicted boxes. The
      #   latter behaviour is necessary to ensure that the number of matches per
      #   datapoint is the same for all datapoints and shapes are static and jit
      #   compatible.
      padded_shapes = {
          "inputs": [max_size, max_size, 3],
          "padding_mask": [max_size, max_size],
          "label": {
              "area": [
                  max_boxes,
              ],
              "boxes": [max_boxes, 4],
              "objects/id": [
                  max_boxes,
              ],
              "is_crowd": [
                  max_boxes,
              ],
              "labels": [
                  max_boxes,
              ],
              "image/id": [],
              "orig_size": [
                  2,
              ],
              "size": [
                  2,
              ],
          },
      }
  
      if train:
          # First repeat then batch.
          ds = ds.shuffle(shuffle_buffer_size, seed=shuffle_seed)
          ds = ds.repeat()
          # Augmentation should be done after repeat for true randomness.
          ds = ds.map(
              preprocess_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          ds = ds.padded_batch(
              batch_size,
              padded_shapes=padded_shapes,
              drop_remainder=True,
          )
  
      else:
          ds = ds.map(
              preprocess_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          # First batch then repeat.
          ds = ds.padded_batch(
              batch_size,
              padded_shapes=padded_shapes,
              drop_remainder=False,
          )
          ds = ds.repeat()
  
      ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
      return ds, builder.info
  
file_path: scenic/scenic/projects/baselines/detr/input_pipeline_detection.py
imports: import functools
  from typing import Optional
  from absl import logging
  from flax import jax_utils
  import jax
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.dataset_lib.coco_dataset import coco_utils
  from scenic.projects.baselines.detr import transforms
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: coco_load_split_from_tfds(batch_size) and returns: ds,  builder