code: def mesh_inference_dataset_fn(
      mixture_or_task_name,
      sequence_length,
      dataset_split,
      shuffle=False,
      seed=None,
      vocabulary=None,
      num_inference_examples=-1,
      use_cached=False,
      priming_sequence_length=None,
  ):
      """Returns all tf.data.Datasets for LM inference on a given mixture.
  
      For Tasks without inputs (such as language modeling), the first
      `priming_sequence_length` tokens in the target are used as the "inputs" for
      inference.
  
      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. If set to None, packing and padding will be
          disabled.
        dataset_split: string, which split of the dataset to load. NOTE, this
          function does NOT receive the split specified in utils.run. It needs to be
          specified separately.
        shuffle: Whether or not to shuffle dataset.
        seed: tf.int64 scalar tf.Tensor (or None). Used as shuffle seed for tf.data.
        vocabulary: unused argument, maintains compatibility with other dataaset_fns
        num_inference_examples: maximum number of examples per task to do inference
          on. If None or less than 0, use all examples.
        use_cached: bool, whether to load the cached version of this dataset.
          evals but should not be used for iterative decoding.
        priming_sequence_length: If the Task only has "targets", select the first
          this many tokens from each target sequence to use as "inputs". This is
          useful for decoder-only language models where you would like to use a
          portion of the targets as a priming sequence for generation.
  
      Returns:
        A list of mesh_tensorflow.transformer.dataset.EvalDataset tuples.
      """
      del vocabulary
      mixture_or_task = seqio.get_mixture_or_task(mixture_or_task_name)
  
      def _split_targets_for_primed_inference(ex):
          ex["inputs"] = ex["targets"][:priming_sequence_length]
          ex["targets"] = ex["targets"][priming_sequence_length:]
          ex["inputs"] = tf.pad(
              ex["inputs"],
              [
                  [
                      0,
                      priming_sequence_length
                      - tf.shape(ex["inputs"])[0],
                  ]
              ],
              "CONSTANT",
          )
          ex["inputs"] = tf.reshape(
              ex["inputs"], shape=(priming_sequence_length,)
          )
          return ex
  
      def _prepare_for_unprimed_inference(ex):
          ex["inputs"] = tf.constant([], dtype=tf.int64)
          return ex
  
      def _get_dataset_for_single_task(task, sequence_length):
          """Get a tensorflow.data.Dataset for the provided task."""
  
          ds = task.get_dataset(
              sequence_length,
              split=dataset_split,
              use_cached=use_cached,
              shuffle=shuffle,
              seed=seed,
          )
          if "inputs" not in ds.element_spec:
              if (
                  not priming_sequence_length
                  or priming_sequence_length <= 0
              ):
                  logging.warning(
                      "Priming sequence length not specified so priming "
                      "with the empty string."
                  )
                  ds = ds.map(_prepare_for_unprimed_inference)
              else:
                  logging.info(
                      "Using the first %d tokens of each target as input.",
                      priming_sequence_length,
                  )
                  ds = ds.map(_split_targets_for_primed_inference)
          elif priming_sequence_length is not None:
              raise ValueError(
                  "Setting a priming sequence length only makes sense for decoder-only "
                  "Tasks, which have `targets` but no `inputs`."
              )
  
          eos_keys = set(
              k
              for k, f in mixture_or_task.output_features.items()
              if f.add_eos
          )
  
          logging.info(
              "Padding '%s' with sequence lengths: %s",
              task.name,
              sequence_length,
          )
          ds = transformer_dataset.pack_or_pad(
              ds,
              sequence_length,
              pack=False,
              feature_keys=tuple(task.output_features),
              ensure_eos=eos_keys,
          )
  
          if (
              num_inference_examples is not None
              and num_inference_examples >= 0
          ):
              ds = ds.take(num_inference_examples)
  
          return ds
  
      outputs = []
  
      for task in seqio.get_subtasks(mixture_or_task):
          if dataset_split not in task.splits:
              logging.info(
                  "Task %s has no '%s' split, skipping inference.",
                  task.name,
                  dataset_split,
              )
              continue
  
          outputs.append(
              transformer_dataset.EvalDataset(
                  task.name,
                  functools.partial(
                      _get_dataset_for_single_task,
                      task=task,
                      sequence_length=sequence_length,
                  ),
                  task.postprocess_fn,
                  task.metric_fns,
              )
          )
  
      if not outputs:
          logging.warning(
              "No %s data found for %s.",
              dataset_split,
              mixture_or_task_name,
          )
  
      return outputs
  
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 2 for calling a global function with example usage: mesh_inference_dataset_fn(mixture_or_task_name, sequence_length, dataset_split, shuffle, seed, vocabulary, num_inference_examples, use_cached, priming_sequence_length) and returns: outputs, ex, ex, ds