code: def testNonTrivialMeanMatchesGP(self, num_tasks):
      # Check that the independent kernel parameterization matches using a
      # single-task GP.
  
      # 5x5 grid of index points in R^2 and flatten to 9x2
      index_points = np.linspace(-4.0, 4.0, 3, dtype=np.float32)
      index_points = np.stack(
          np.meshgrid(index_points, index_points), axis=-1
      )
      index_points = np.reshape(index_points, [-1, 2])
      # ==> shape = [9, 2]
  
      amplitude = np.float32(0.5)
      length_scale = np.float32(2.0)
      kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          amplitude, length_scale
      )
      observation_noise_variance = np.float32(1e-2)
      multi_task_kernel = multitask_kernel.Independent(
          num_tasks=num_tasks, base_kernel=kernel
      )
  
      observations = (
          np.linspace(-20.0, 20.0, 9 * num_tasks)
          .reshape(9, num_tasks)
          .astype(np.float32)
      )
  
      test_points = np.random.uniform(-1.0, 1.0, [10, 2]).astype(
          np.float32
      )
      test_observations = np.random.uniform(
          -20.0, 20.0, [10, num_tasks]
      ).astype(np.float32)
  
      # Constant mean per task.
      mean_fn = lambda x: tf.linspace(1.0, 3.0, num_tasks)
  
      mtgp = mtgprm_lib.MultiTaskGaussianProcessRegressionModel(
          multi_task_kernel,
          observation_index_points=index_points,
          index_points=test_points,
          observations=observations,
          observation_noise_variance=observation_noise_variance,
          mean_fn=mean_fn,
          validate_args=True,
      )
  
      # For the single task GP, we move the task dimension to the front of the
      # batch shape.
      gp = gprm.GaussianProcessRegressionModel(
          kernel,
          observation_index_points=index_points,
          index_points=test_points,
          observations=tf.linalg.matrix_transpose(observations),
          observation_noise_variance=observation_noise_variance,
          mean_fn=lambda x: tf.linspace(1.0, 3.0, num_tasks)[
              ..., tf.newaxis
          ],
          validate_args=True,
      )
      # Print batch of covariance matrices.
      multitask_log_prob = mtgp.log_prob(test_observations)
      single_task_log_prob = tf.reduce_sum(
          gp.log_prob(tf.linalg.matrix_transpose(test_observations)),
          axis=0,
      )
      self.assertAllClose(
          self.evaluate(single_task_log_prob),
          self.evaluate(multitask_log_prob),
          rtol=4e-3,
      )
  
      multi_task_mean_ = self.evaluate(mtgp.mean())
      # Reshape so that task dimension is last.
      single_task_mean_ = np.swapaxes(self.evaluate(gp.mean()), -1, -2)
      self.assertAllClose(
          single_task_mean_, multi_task_mean_, rtol=1e-5
      )
  
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 7 for calling the MultiTaskGaussianProcessRegressionModelTest obj's underlying member method code with example usage: obj.testNonTrivialMeanMatchesGP(num_tasks) without return types