code: def _build_pipeline(
      config: ml_collections.ConfigDict,
      batch_size: int,
      rng: Any,
      shuffle: bool = False,
  ) -> tf.data.Dataset:
      """Build a tf.data.Dataset pipeline using clu.deterministic_data.
  
      The pipeline has the following steps:
  
      0. All used datasets should have text label fields. Convert labels to text in
         their decoders.
  
      1. Datasets are merged with tf.data.Dataset.sample_from_datasets.
  
      2. Optionally, mosaics are created.
  
      3. Prompts are added to text labels.
  
      4. Per-image label spaces are generated from text labels.
  
      5. Text queries are tokenized.
  
      6. Features are converted to Scenic format.
  
      Args:
        config: Dataset configurations.
        batch_size: Total batch size (sum for all devices).
        rng: Random seed (JAX format).
        shuffle: Whether to shuffle.
  
      Returns:
        tf.data.Dataset after preprocessing, merging, mosaicing, and batching.
      """
  
      builders = [
          tfds.builder(name, data_dir=config.data_dirs[0].get(name))
          for name in config.tfds_names
      ]
  
      decoder_kwarg_list = config.get(
          "decoder_kwarg_list", [{}] * len(config.tfds_names)
      )
      mosaic_sizes = config.get("mosaic_sizes", (1,))
      mosaic_probs = config.get("mosaic_probs", (1.0,))
      mosaic_datasets = []
      for mosaic_size in mosaic_sizes:
          rng, ds_rng = jax.random.split(rng)
          merged_dataset = _get_merged_dataset(
              builders=builders,
              splits=config.splits,
              dataset_probs=config.dataset_probs,
              decoder_kwarg_list=decoder_kwarg_list,
              preproc_spec=config.preproc_spec,
              mosaic_size=mosaic_size,
              rng=ds_rng,
              shuffle=shuffle,
              shuffle_buffer_size=config.shuffle_buffer_size,
              cache=config.get("cache", False),
          )
  
          # Build mosaics:
          if mosaic_size == 1:  # I.e. no mosaic.
              merged_dataset = (
                  merged_dataset.unbatch()
              )  # Remove mosaic dimension.
          else:
              merged_dataset = merged_dataset.map(
                  mosaic.CreateMosaic(mosaic_size),
                  num_parallel_calls=NUM_PARALLEL_CALLS,
              )
  
          mosaic_datasets.append(merged_dataset)
  
      # Merge mosaic sizes:
      final_dataset = tf.data.Dataset.sample_from_datasets(
          mosaic_datasets, weights=mosaic_probs, seed=rng[0]
      )
  
      # Apply post-mosaic processing:
      post_mosaic_processing = _get_post_mosaic_process_fn(
          config.preproc_spec
      )
      post_mosaic_ops = list(post_mosaic_processing.ops)
  
      post_mosaic_ops.append(
          label_ops.ConvertToScenic(
              input_range=config.input_range,
          )
      )
  
      post_mosaic_processing.ops = post_mosaic_ops
      final_dataset = final_dataset.map(
          post_mosaic_processing, num_parallel_calls=NUM_PARALLEL_CALLS
      )
  
      # Batch to the desired output batch size:
      batch_dims = [
          jax.local_device_count(),
          batch_size // jax.local_device_count(),
      ]
      for batch_size in reversed(batch_dims):
          final_dataset = final_dataset.batch(
              batch_size,
              drop_remainder=True,
              num_parallel_calls=NUM_PARALLEL_CALLS,
          )
  
      return final_dataset.prefetch(tf.data.AUTOTUNE)
  
imports: import dataclasses
  import functools
  from typing import Any, Dict, Optional, Sequence
  from clu import deterministic_data
  from clu import preprocess_spec
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.owl_vit.preprocessing import image_ops
  from scenic.projects.owl_vit.preprocessing import label_ops
  from scenic.projects.owl_vit.preprocessing import mosaic
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 5 for calling a global function with example usage: _build_pipeline(config, batch_size, rng, shuffle) and returns: final_dataset