code: def test_get_multi_dataset(self, pmf, num_datasets, sample_size):
      with self.test_session() as sess:
          datasets = [
              tf.data.Dataset.from_tensors(i)
              for i in range(num_datasets)
          ]
          multi_dataset = multi_problem_v2.get_multi_dataset(
              datasets, pmf
          )
          multi_dataset = multi_dataset.batch(sample_size)
          iterator = multi_dataset.make_initializable_iterator()
          sess.run(iterator.initializer)
          sample_pmf = tf.reduce_mean(
              tf.one_hot(iterator.get_next(), num_datasets), 0
          )
          if pmf is None:
              pmf = np.array(
                  [1.0 / num_datasets] * num_datasets, np.float32
              )
          self.assertAllClose(
              sess.run(sample_pmf), pmf, rtol=0.1, atol=0.1
          )
  
file_path: tensor2tensor/tensor2tensor/data_generators/multi_problem_v2_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  from tensor2tensor.data_generators import multi_problem_v2
  from tensor2tensor.data_generators import problem
  import tensorflow.compat.v1
label_desc: Use image node_id 7 for calling the MultiProblemV2Test obj's underlying member method code with example usage: obj.test_get_multi_dataset(pmf, num_datasets, sample_size) without return types