import tensorflow as tf

    p.task_params.Define('a', p0, '')
    p.task_params.Define('b', p1, '')

    return p

  def _testSampleTaskHelper(self, p):
    model = p.cls(p)

    task_to_id = {model.children['a']: 'a', model.children['b']: 'b'}
    task_counts = {'a': 0, 'b': 0}

    # initialize tensorflow graph and global step
    with self.session() as sess:
      tf.global_variables_initializer().run()
      global_step = sess.run(model.global_step)
      for _ in range(100):
        task = model.SampleTask(global_step)
        task_counts[task_to_id[task]] += 1

      self.assertEqual(task_counts['a'], 83)
      self.assertEqual(task_counts['b'], 17)

  def testSampleTaskSpecifiedWithoutScheduler(self):
    """Expected distribution: 'a': 0.8 , 'b': 0.2."""
    p = self._setUpTestSampleTask()

