code: def mesh_train_dataset_fn(
      mixture_or_task_name,
      sequence_length,
      vocabulary=None,
      dataset_split=tfds.Split.TRAIN,
      shuffle=True,
      seed=None,
      use_cached=False,
      pack=True,
  ):
      """Returns the tf.data.Dataset for training on a given mixture.
  
      This uses the format required for utils.run's `train_dataset_fn` argument in
      the Mesh TF transformer standalone.
  
      Args:
        mixture_or_task_name: string, an identifier for a Mixture or Task in the
          appropriate registry. Must be specified via gin.
        sequence_length: dict mapping feature key to the int length for that feature
          the max sequence length.
        vocabulary: unused argument, maintains compatibility with other dataset_fns.
        dataset_split: string, which split of the dataset to load. In most cases
          this should be "train".
        shuffle: Whether or not to shuffle dataset.
        seed: tf.int64 scalar tf.Tensor (or None). Used for both the global seed and
          shuffle seed for tf.data
        use_cached: bool, whether to load the cached version of this dataset.
        pack: bool, whether to pack the dataset.
  
      Returns:
        A tf.data.Dataset of preprocessed, tokenized, and batched examples.
      """
      del vocabulary
      mixture_or_task = seqio.get_mixture_or_task(mixture_or_task_name)
  
      ds = mixture_or_task.get_dataset(
          sequence_length,
          split=dataset_split,
          use_cached=use_cached,
          shuffle=shuffle,
          num_epochs=None,
          seed=seed,
      )
  
      # Select just the output features which are present in the dataset.
      feature_keys = tuple(
          k
          for k in mixture_or_task.output_features
          if k in tf.data.get_output_shapes(ds)
      )
  
      # Filtering feature keys is done in pack_or_pad function. However, when
      # packing is turned off, input_features aren't filtered leading to training
      # problems due to strings showing up in the input example. Filtering features
      # ensures that we don't rely on pack_or_pad to filter features for training.
      def _filter_features(ex):
          return {k: ex[k] for k in feature_keys}
  
      ds = ds.map(
          _filter_features,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
  
      eos_keys = set(
          k
          for k, f in mixture_or_task.output_features.items()
          if f.add_eos
      )
      ds = transformer_dataset.pack_or_pad(
          ds,
          sequence_length,
          pack=pack,
          feature_keys=feature_keys,
          ensure_eos=eos_keys,
      )
      return ds
  
file_path: text-to-text-transfer-transformer/t5/models/mesh_transformer.py
imports: import functools
  from absl import logging
  import gin
  import mesh_tensorflow.transformer.dataset
  import seqio
  from t5.models import utils
  import tensorflow.compat.v1
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: mesh_train_dataset_fn(mixture_or_task_name, sequence_length, vocabulary, dataset_split, shuffle, seed, use_cached, pack) and returns: ds, unknown