code: def test_generate_batch_output_shape(
      self, x_dim, prior_sigma, sigma_min, sigma_max, n, batch_size
  ):
      partition_generator = DummyPartitionGenerator()
      mog_generator = mog.MOGGenerator(
          partition_generator=partition_generator,
          x_dim=x_dim,
          prior_sigma=prior_sigma,
          sigma_min=sigma_min,
          sigma_max=sigma_max,
      )
      data, labels = mog_generator.generate_batch(
          n=n, batch_size=batch_size
      )
      np.testing.assert_array_equal(
          data.shape, np.array([batch_size, n, x_dim])
      )
      np.testing.assert_array_equal(
          labels.shape, np.array([batch_size, n])
      )
  
file_path: neural-structured-learning/research/neural_clustering/data_generators/mog_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from neural_clustering.data_generators import mog
  import numpy
label_desc: Use image node_id 1 for calling the MogTest obj's underlying member method code with example usage: obj.test_generate_batch_output_shape(x_dim, prior_sigma, sigma_min, sigma_max, n, batch_size) without return types