code: def dataset(self, mode, hparams=None, global_step=None, **kwargs):
      
      datasets = [p.dataset(mode, **kwargs) for p in self.problems]
      datasets = [
          d.map(lambda x, i=j: self.normalize_example(  # pylint: disable=g-long-lambda
              dict(x, problem_id=tf.constant([i])), hparams))
          for j, d in enumerate(datasets)  # Tag examples with a problem_id.
      ]
      if mode is problem.DatasetSplit.TRAIN:
        if global_step is None:
          global_step = tf.train.get_or_create_global_step()
        pmf = get_schedule_distribution(self.schedule, global_step)
        return get_multi_dataset(datasets, pmf)
      elif self.only_eval_first_problem:
        return datasets[0]
      else:
        datasets = [d.repeat() for d in datasets]
        return tf.data.Dataset.zip(tuple(datasets)).flat_map(
            lambda *x: functools.reduce(  # pylint: disable=g-long-lambda
                tf.data.Dataset.concatenate,
                map(tf.data.Dataset.from_tensors, x)))
docs: Returns a dataset containing examples from multiple problems.
  
      Args:
        mode: A member of problem.DatasetSplit.
        hparams: A tf.HParams object, the model hparams.
        global_step: A scalar tensor used to compute the sampling distribution.
          If global_step is None, we call tf.train.get_or_create_global_step by
          default.
        **kwargs: Keywords for problem.Problem.Dataset.
  
      Returns:
        A dataset containing examples from multiple problems.