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.

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)