code: def test_max_value_distribution(self):
      # Test that we can create the max-value distribution
      # and it has the right shape.
  
      index_points = np.linspace(-4.0, 4.0, 5, dtype=self.dtype)
      index_points = np.stack(
          np.meshgrid(index_points, index_points), axis=-1
      )
      index_points = np.reshape(index_points, [-1, 2])
      index_points = np.stack([index_points] * 6)
  
      observation_index_points = np.linspace(
          4.0, 10.0, 4, dtype=self.dtype
      )
      observation_index_points = np.stack(
          np.meshgrid(
              observation_index_points, observation_index_points
          ),
          axis=-1,
      )
      observation_index_points = np.reshape(
          observation_index_points, [-1, 2]
      )
  
      observations = np.linspace(2.0, 10.0, 16, dtype=self.dtype)
  
      # Kernel with batch_shape [2, 4, 1, 1]
      amplitude = np.array([1.0, 2.0], self.dtype).reshape([2, 1, 1, 1])
      length_scale = np.array([1.0, 2.0, 3.0, 4.0], self.dtype).reshape(
          [1, 4, 1, 1]
      )
  
      predictive_distribution = gprm.GaussianProcessRegressionModel(
          kernel=exponentiated_quadratic.ExponentiatedQuadratic(
              amplitude=amplitude, length_scale=length_scale
          ),
          observations=observations,
          observation_index_points=observation_index_points,
          index_points=index_points,
      )
  
      max_value_distribution = fit_max_value_distribution(
          predictive_distribution=predictive_distribution,
          observations=observations,
          num_grid_points=100,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
  
      actual_mes = (
          max_value_entropy_search.GaussianProcessMaxValueEntropySearch(
              predictive_distribution=predictive_distribution,
              observations=observations,
              num_max_value_samples=1,
              seed=test_util.test_seed(),
          )()
      )
  
      # Ensure that we properly take in to account the batch shape from various
      # sources.
      self.assertEqual(max_value_distribution.batch_shape, [2, 4, 1, 6])
      # Ensure that the max-value distribution is scalar.
      self.assertEqual(max_value_distribution.event_shape, [])
  
      # Test MES shape.
      shape = [2, 4, 1, 6, 25]
      self.assertAllEqual(actual_mes.shape, shape)
      self.assertDTypeEqual(actual_mes, self.dtype)
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gaussian_process_regression_model
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.bayesopt.acquisition import max_value_entropy_search
  from tensorflow_probability.python.experimental.bayesopt.acquisition.max_value_entropy_search import fit_max_value_distribution
  from tensorflow_probability.python.experimental.bayesopt.acquisition.max_value_entropy_search import inverse_mills_ratio
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
label_desc: Use image node_id 2 for calling the _MaxValueEntropySearchTest obj's underlying member method code with example usage: obj.test_max_value_distribution() without return types