code: def get_dataset(
      *,
      batch_size: int,
      eval_batch_size: int,
      num_shards: int,
      dtype_str: str = "float32",
      shuffle_seed: int = 0,
      rng: Optional[jnp.ndarray] = None,
      dataset_configs: Optional[Mapping[str, Any]] = None,
      dataset_service_address: Optional[str] = None,
  ):
      """Returns generators for COCO object detection 2017 train & validation set.
  
      Args:
        batch_size: int; Determines the train batch size.
        eval_batch_size: int; Determines the evaluation batch size.
        num_shards: int;  Number of shards --> batch shape: [num_shards, bs, ...].
        dtype_str: Data type of the image. Only 'float32' is currently supported.
        shuffle_seed: int; Seed for shuffling the training data.
        rng: JAX rng key, which can be used for augmentation, shuffling, etc.
        dataset_configs: dict; Dataset specific configurations. Must be empty.
        dataset_service_address: If set, will distribute the training dataset using
          the given tf.data service at the given address.
  
      Returns:
        A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
        a test_iter, and a dict of meta_data.
      """
      del rng
      assert (
          dtype_str == "float32"
      ), f"coco_detr_dataset invoked with unsupported dtype_str: {dtype_str}"
      del dtype_str
  
      dataset_configs = dataset_configs or {}
  
      max_size = dataset_configs.get("max_size", 1333)
      max_boxes = dataset_configs.get("max_boxes", 100)
  
      train_preprocess_fn = make_coco_transforms("train", max_size)
      valid_max_size = dataset_configs.get("valid_max_size", max_size)
      eval_preprocess_fn = make_coco_transforms(
          "validation", valid_max_size
      )
  
      decode_fn = decode_coco_deformable_detr_example
  
      train_ds, train_ds_info = coco_load_split_from_tfds(
          batch_size,
          train=True,
          preprocess_fn=train_preprocess_fn,
          decode_fn=decode_fn,
          shuffle_buffer_size=dataset_configs.get(
              "shuffle_buffer_size", 1000
          ),
          max_size=max_size,
          max_boxes=max_boxes,
          shuffle_seed=shuffle_seed,
          data_dir=dataset_configs.get("data_dir"),
      )
  
      if dataset_service_address:
          if shuffle_seed is not None:
              raise ValueError(
                  "Using dataset service with a random seed causes each "
                  "worker to produce exactly the same data. Add "
                  "config.shuffle_seed = None to your config if you "
                  "want to run with dataset service."
              )
          logging.info(
              "Using the tf.data service at %s", dataset_service_address
          )
          train_ds = dataset_utils.distribute(
              train_ds, dataset_service_address
          )
  
      eval_ds, _ = coco_load_split_from_tfds(
          eval_batch_size,
          train=False,
          preprocess_fn=eval_preprocess_fn,
          max_size=valid_max_size,
          max_boxes=max_boxes,
          decode_fn=decode_fn,
          data_dir=dataset_configs.get("data_dir"),
      )
  
      # Labels take on values 1-90. We set 0 to be padded objects.
      num_classes = (
          train_ds_info.features["objects"]["label"].num_classes + 1
      )
  
      maybe_pad_batches_train = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=True,
          batch_size=batch_size,
      )
      maybe_pad_batches_eval = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
      )
      shard_batches = functools.partial(
          dataset_utils.shard, n_devices=num_shards
      )
  
      train_iter = iter(train_ds)
      train_iter = map(dataset_utils.tf_to_numpy, train_iter)
      train_iter = map(maybe_pad_batches_train, train_iter)
      train_iter = map(shard_batches, train_iter)
      if dataset_configs.get("prefetch_to_device"):
          # Async bind batch to device which speeds up training.
          train_iter = jax_utils.prefetch_to_device(
              train_iter, dataset_configs.get("prefetch_to_device")
          )
  
      eval_iter = iter(eval_ds)
      eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
      eval_iter = map(maybe_pad_batches_eval, eval_iter)
      eval_iter = map(shard_batches, eval_iter)
  
      meta_data = {
          "num_classes": num_classes,
          "input_shape": [-1, max_size, max_size, 3],
          "num_train_examples": dataset_utils.get_num_examples(
              "coco/2017",
              "train",
              data_dir=dataset_configs.get("data_dir"),
          ),
          "num_eval_examples": dataset_utils.get_num_examples(
              "coco/2017",
              "validation",
              data_dir=dataset_configs.get("data_dir"),
          ),
          "input_dtype": jnp.float32,
          "target_is_onehot": False,
          "label_to_name": coco_utils.get_label_map("ref_coco"),
      }
      return dataset_utils.Dataset(
          train_iter, eval_iter, None, meta_data
      )
  
imports: import functools
  from typing import Any, Callable, Mapping, 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 6 for calling a global function with example usage: get_dataset() and returns: dataset_utils