code: def test_max_posterior_sampling(self):
      batch_shapes = (torch.Size(), torch.Size([3]), torch.Size([3, 2]))
      dtypes = (torch.float, torch.double)
      for batch_shape, dtype, N, num_samples, d in itertools.product(
          batch_shapes, dtypes, (5, 6), (1, 2), (1, 2)
      ):
          tkwargs = {"device": self.device, "dtype": dtype}
          # X is `batch_shape x N x d` = batch_shape x N x 1.
          X = torch.randn(*batch_shape, N, d, **tkwargs)
          # the event shape is `num_samples x batch_shape x N x m`
          psamples = torch.zeros(
              num_samples, *batch_shape, N, 1, **tkwargs
          )
          psamples[..., 0, :] = 1.0
  
          # IdentityMCObjective, with replacement
          with mock.patch.object(
              MockPosterior, "rsample", return_value=psamples
          ):
              mp = MockPosterior(None)
              with mock.patch.object(
                  MockModel, "posterior", return_value=mp
              ):
                  mm = MockModel(None)
                  MPS = MaxPosteriorSampling(mm)
                  s = MPS(X, num_samples=num_samples)
                  self.assertTrue(
                      torch.equal(s, X[..., [0] * num_samples, :])
                  )
  
          # ScalarizedPosteriorTransform w/ replacement
          with mock.patch.object(
              MockPosterior, "rsample", return_value=psamples
          ):
              mp = MockPosterior(None)
              with mock.patch.object(
                  MockModel, "posterior", return_value=mp
              ):
                  mm = MockModel(None)
                  with mock.patch.object(
                      ScalarizedPosteriorTransform,
                      "forward",
                      return_value=mp,
                  ):
                      post_tf = ScalarizedPosteriorTransform(
                          torch.rand(2, **tkwargs)
                      )
                      MPS = MaxPosteriorSampling(
                          mm, posterior_transform=post_tf
                      )
                      s = MPS(X, num_samples=num_samples)
                      self.assertTrue(
                          torch.equal(s, X[..., [0] * num_samples, :])
                      )
  
          # without replacement
          psamples[..., 1, 0] = 1e-6
          with mock.patch.object(
              MockPosterior, "rsample", return_value=psamples
          ):
              mp = MockPosterior(None)
              with mock.patch.object(
                  MockModel, "posterior", return_value=mp
              ):
                  mm = MockModel(None)
                  MPS = MaxPosteriorSampling(mm, replacement=False)
                  if len(batch_shape) > 1:
                      with self.assertRaises(NotImplementedError):
                          MPS(X, num_samples=num_samples)
                  else:
                      s = MPS(X, num_samples=num_samples)
                      # order is not guaranteed, need to sort
                      self.assertTrue(
                          torch.equal(
                              torch.sort(s, dim=-2).values,
                              torch.sort(
                                  X[..., :num_samples, :], dim=-2
                              ).values,
                          )
                      )
  
imports: from __future__ import annotations
  import itertools
  from unittest import mock
  import torch
  from botorch.acquisition.analytic import PosteriorMean
  from botorch.acquisition.objective import IdentityMCObjective, LinearMCObjective, ScalarizedPosteriorTransform
  from botorch.generation.sampling import BoltzmannSampling, ConstrainedMaxPosteriorSampling, MaxPosteriorSampling, SamplingStrategy
  from botorch.models import SingleTaskGP
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
label_desc: Use image node_id 2 for calling the TestMaxPosteriorSampling obj's underlying member method code with example usage: obj.test_max_posterior_sampling() without return types