code: def _get_merged_dataset(
      builders: Sequence[tfds.core.DatasetBuilder],
      splits: Sequence[str],
      dataset_probs: Sequence[float],
      decoder_kwarg_list: Sequence[Dict[str, Any]],
      preproc_spec: str,
      mosaic_size: int,
      rng: Any,
      shuffle: bool = False,
      shuffle_buffer_size: int = 10_000,
      cache: bool = False,
  ) -> tf.data.Dataset:
      """Creates datasets from builders, applies preprocessing, and merges them.
  
      Args:
        builders: List of TFDS dataset builders.
        splits: List containing the split to use for each builder.
        dataset_probs: Sampling probabilities for each dataset.
        decoder_kwarg_list: Kwargs to pass to the decoder.
        preproc_spec: Preprocessing specification string.
        mosaic_size: Number of tiles along each mosaic edge. For no mosaicing, set
          to 1.
        rng: Random seed (JAX format).
        shuffle: Whether to shuffle.
        shuffle_buffer_size: Shuffle buffer size (ignored if shuffle is False).
        cache: Whether to cache the datasets.
  
      Returns:
        tf.data.Dataset with pre-mosaic preprocessing applied.
      """
  
      datasets_to_merge = []
      for builder, split, decoder_kwargs in zip(
          builders, splits, decoder_kwarg_list
      ):
          pre_mosaic_processing = _get_pre_mosaic_process_fn(
              builder, decoder_kwargs, preproc_spec, mosaic_size
          )
          rng, ds_rng = jax.random.split(rng)
          datasets_to_merge.append(
              _get_single_tfds_dataset(
                  builder,
                  split,
                  batch_size=mosaic_size**2,
                  preprocess_fn=pre_mosaic_processing,
                  rng=ds_rng,
                  shuffle=shuffle,
                  shuffle_buffer_size=shuffle_buffer_size,
                  cache=cache,
              )
          )
  
      try:
          return tf.data.Dataset.sample_from_datasets(
              datasets_to_merge,
              weights=dataset_probs,
              seed=None if rng is None else rng[0],
          )
      except TypeError as e:
          # There is a mismatch between the datasets to be merged. If it's a simple
          # difference in the feature keys, provide a nicer message than tf.data:
          expected_keys = set(datasets_to_merge[0].element_spec.keys())
          for dataset in datasets_to_merge[1:]:
              actual_keys = set(dataset.element_spec.keys())
              if actual_keys != expected_keys:
                  raise TypeError(
                      "Datasets to be merged must have the same structure, but had keys "
                      f"\n\n{expected_keys}\n\n and \n\n{actual_keys}\n\n"
                      f"Difference:\n{expected_keys.symmetric_difference(actual_keys)}"
                  ) from e
          # If the difference is not in the feature keys, raise the original error:
          raise e
  
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 4 for calling a global function with example usage: _get_merged_dataset(builders, splits, dataset_probs, decoder_kwarg_list, preproc_spec, mosaic_size, rng, shuffle, shuffle_buffer_size, cache) and returns: tf