code: def get_multi_dataset(datasets, pmf=None):
      """Returns a Dataset that samples records from one or more Datasets.
  
      Args:
        datasets: A list of one or more Dataset objects to sample from.
        pmf: A tensor of shape [len(datasets)], the probabilities to sample each
          dataset with. This tensor is often constructed with the global_step. If
          this is None, we sample from the datasets uniformly at random.
  
      Returns:
        A Dataset object containing records from multiple datasets. Note that
        because this dataset iterates through other datasets it is stateful, thus
        you will need to call make_initializable_iterator instead of
        make_one_shot_iterator.
      """
      pmf = (
          tf.fill([len(datasets)], 1.0 / len(datasets))
          if pmf is None
          else pmf
      )
      samplers = [
          d.repeat().make_one_shot_iterator().get_next for d in datasets
      ]
      sample = lambda _: categorical_case(pmf, samplers)
      return tf.data.Dataset.from_tensors([]).repeat().map(sample)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import numpy
  from tensor2tensor.data_generators import problem
  from tensor2tensor.data_generators import text_problems
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: get_multi_dataset(datasets, pmf) and returns: tf