code: def testShapes(self, num_tasks):
      # 3x3 grid of index points in R^2 and flatten to 9x2
      index_points = np.linspace(-4.0, 4.0, 3, dtype=np.float64)
      index_points = np.stack(
          np.meshgrid(index_points, index_points), axis=-1
      )
      index_points = np.reshape(index_points, [-1, 2])
  
      batched_index_points = np.stack([index_points] * 6)
      # ==> shape = [6, 9, 2]
  
      # ==> shape = [9, 2]
      observations = np.linspace(-20.0, 20.0, num_tasks * 9).reshape(
          9, num_tasks
      )
  
      test_index_points = np.random.uniform(-6.0, 6.0, [5, 2])
      # ==> shape = [3, 1, 5, 2]
  
      # Kernel with batch_shape [2, 4, 3, 1, 1]
      amplitude = np.array([1.0, 2.0], np.float64).reshape([2, 1, 1, 1])
      length_scale = np.array([1.0, 2.0, 3.0, 4.0], np.float64).reshape(
          [1, 4, 1, 1]
      )
      observation_noise_variance = np.array(
          [1e-5, 1e-6, 1e-5], np.float64
      ).reshape([1, 1, 3, 1])
      kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          amplitude, length_scale
      )
      multi_task_kernel = multitask_kernel.Independent(
          num_tasks=num_tasks, base_kernel=kernel
      )
      gp = mtgprm_lib.MultiTaskGaussianProcessRegressionModel(
          multi_task_kernel,
          observation_index_points=batched_index_points,
          observations=observations,
          index_points=test_index_points,
          observation_noise_variance=observation_noise_variance,
          predictive_noise_variance=0.0,
          validate_args=True,
      )
  
      batch_shape = [2, 4, 3, 6]
      event_shape = [5, num_tasks]
      sample_shape = [5, 3]
  
      samples = gp.sample(sample_shape, seed=test_util.test_seed())
  
      self.assertAllEqual(gp.batch_shape, batch_shape)
      self.assertAllEqual(gp.event_shape, event_shape)
      self.assertAllEqual(
          self.evaluate(gp.batch_shape_tensor()), batch_shape
      )
      self.assertAllEqual(
          self.evaluate(gp.event_shape_tensor()), event_shape
      )
      self.assertAllEqual(
          self.evaluate(samples).shape,
          sample_shape + batch_shape + event_shape,
      )
      self.assertAllEqual(
          self.evaluate(gp.log_prob(samples)).shape,
          sample_shape + batch_shape,
      )
      self.assertAllEqual(
          self.evaluate(tf.shape(gp.mean())), batch_shape + event_shape
      )
  
imports: from unittest import mock
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gaussian_process_regression_model
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.distributions import multitask_gaussian_process_regression_model
  from tensorflow_probability.python.experimental.psd_kernels import multitask_kernel
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.math.psd_kernels.internal import test_util
label_desc: Use image node_id 2 for calling the MultiTaskGaussianProcessRegressionModelTest obj's underlying member method code with example usage: obj.testShapes(num_tasks) without return types