code: def testCustomMarginalFn(self):
      def test_marginal_fn(
          loc,
          covariance,
          validate_args=False,
          allow_nan_stats=False,
          name="custom_marginal",
      ):
          return mvn_diag.MultivariateNormalDiag(
              loc=loc,
              scale_diag=tf.math.sqrt(tf.linalg.diag_part(covariance)),
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              name=name,
          )
  
      index_points = np.expand_dims(
          np.random.uniform(-1.0, 1.0, 10), -1
      )
  
      gp = gaussian_process.GaussianProcess(
          kernel=psd_kernels.ExponentiatedQuadratic(),
          index_points=index_points,
          marginal_fn=test_marginal_fn,
          validate_args=True,
      )
  
      self.assertAllClose(
          np.eye(10), gp.get_marginal_distribution().covariance()
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gaussian_process
  from tensorflow_probability.python.distributions import gaussian_process_regression_model
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math import psd_kernels
  from tensorflow_probability.python.math.psd_kernels.internal import test_util
label_desc: Use image node_id 8 for calling the _GaussianProcessTest obj's underlying member method code with example usage: obj.testCustomMarginalFn() and returns: mvn_diag