code: def get_multi_dataset(datasets, pmf=None):
    
    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)
docs: 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.