def mix_data(example):
        """Function to mix the different datasets according to a schedule."""
        del example
        # This block computes the probability of mixing the primary task with
        # the secondary tasks. 0 = only the primary task, 1 = only the secondary
        # tasks.
        if hparams.multiproblem_mixing_schedule == MixingSchedule.EXPONENTIAL:
          prob = get_exp_sched_prob()
          prob = tf.cond(
              tf.equal(tf.floormod(
                  problem_step, tf.cast(5e6, dtype=tf.int64)), 0),
              lambda: tf.Print(prob, [prob], message="Probability"),
              lambda: prob)
        elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
          prob = get_const_sched_prob()
        elif hparams.multiproblem_mixing_schedule == MixingSchedule.PRETRAIN:
          prob = get_pretrain_sched_prob()
        else:
          raise ValueError("Unknown schedule %s" % str(
              hparams.multiproblem_mixing_schedule))
        tf.logging.info("Using the %s schedule to "
                        "train the MultiProblem." % str(
                            hparams.multiproblem_mixing_schedule))
        tf.logging.info("Schedule mixing threshold "
                        "%.2f" % hparams.multiproblem_schedule_threshold)

        # If per-task thresholds are specified, use them.
        thresholds = None
        if hparams.multiproblem_per_task_threshold:
          thresholds = hparams.multiproblem_per_task_threshold.split(",")
          thresholds = [float(t) for t in thresholds]  # Convert to floats.
          thresholds_sum = sum(thresholds)
          tf.logging.info("Per-task thresholds: %s." % str(thresholds))
          thresholds = [t / thresholds_sum for t in thresholds]  # Normalize.
          thresholds = [sum(thresholds[:i+1]) for i in range(len(thresholds))]
          tf.logging.info("Per-task threshold sums: %s." % str(thresholds))
          if len(thresholds) != len(self.task_list):
            tf.logging.warn("Specified %d thresholds but encountered %d tasks."
                            % (len(thresholds), len(self.task_list)))
            thresholds = None

        def sample_task(curr_task, num_tasks_left, randnum):
          """A recursive function to sample a task.

          This function treats the probability as the threshold for the primary
          task and divides the remaining probability mass across the other
          tasks.

          Args:
            curr_task: The index of the task being considered for sampling.
            num_tasks_left: Number of tasks remaining to possibly sample from.
            randnum: The random number used to select the dataset.

          Returns:
            A Tensor representing an example from the task that was sampled
            from.
          """
          if num_tasks_left == 0:
            return get_next_from_dataset(dataset_iterators[curr_task])

          if thresholds is not None:  # Use per-task thresholds if specified.
            prob_sum = thresholds[curr_task]
            return tf.cond(
                randnum < prob_sum,
                lambda: get_next_from_dataset(dataset_iterators[curr_task]),
                lambda: sample_task(curr_task+1, num_tasks_left-1, randnum)
            )

          # When curr_task is 0, the primary task, the new prob is the same as
          # the original probability. `tf.greater` indicates that the primary
          # task receives (1-prob) of the probability mass.
          # Otherwise, `prob` is divided equally amongst all the secondary
          # tasks.
          new_prob = prob - (curr_task * prob / (len(self.task_list)-1))
          return tf.cond(
              tf.greater(randnum, new_prob),
              lambda: get_next_from_dataset(dataset_iterators[curr_task]),
              lambda: sample_task(curr_task+1, num_tasks_left-1, randnum)
          )

        return tf.data.Dataset.from_tensors(
            sample_task(0, len(self.task_list)-1, tf.random_uniform([])))
Function to mix the different datasets according to a schedule.
tensor2tensor/data_generators/multi_problem.py