Added posterior_transform to posterior method in ApproximateGPyTorchModel

botorch/models/approximate_gp.py

@@ -31,10 +31,10 @@
 
 import copy
 import warnings
-
 from typing import Optional, Union
 
 import torch
+from botorch.acquisition.objective import PosteriorTransform
 from botorch.exceptions.warnings import UserInputWarning
 from botorch.models.gpytorch import GPyTorchModel
 from botorch.models.transforms.input import InputTransform
@@ -146,8 +146,16 @@ def train(self, mode: bool = True) -> Self:
  return Module.train(self, mode=mode)
 
  def posterior(
- self, X, output_indices=None, observation_noise=False, *args, **kwargs
+ self,
+ X,
+ output_indices: Optional[list[int]] = None,
+ observation_noise: bool = False,
+ posterior_transform: Optional[PosteriorTransform] = None,
  ) -> GPyTorchPosterior:
+ if output_indices is not None:
+ raise NotImplementedError( # pragma: no cover
+ f"{self.__class__.__name__}.posterior does not support output indices."
+ )
  self.eval() # make sure model is in eval mode
 
  # input transforms are applied at `posterior` in `eval` mode, and at
@@ -161,11 +169,13 @@ def posterior(
  X = X.unsqueeze(-3).repeat(*[1] * (X_ndim - 2), self.num_outputs, 1, 1)
  dist = self.model(X)
  if observation_noise:
- dist = self.likelihood(dist, *args, **kwargs)
+ dist = self.likelihood(dist)
 
  posterior = GPyTorchPosterior(distribution=dist)
  if hasattr(self, "outcome_transform"):
  posterior = self.outcome_transform.untransform_posterior(posterior)
+ if posterior_transform is not None:
+ posterior = posterior_transform(posterior)
  return posterior
 
  def forward(self, X) -> MultivariateNormal:

test/models/test_approximate_gp.py

@@ -8,6 +8,7 @@
 import warnings
 
 import torch
+from botorch.acquisition.objective import ScalarizedPosteriorTransform
 from botorch.exceptions.warnings import UserInputWarning
 from botorch.fit import fit_gpytorch_mll
 from botorch.models.approximate_gp import (
@@ -103,6 +104,16 @@ def test_posterior(self):
  # test batch_shape property
  self.assertEqual(model.batch_shape, tx.shape[:-2])
 
+ # Test that checks if posterior_transform is correctly applied
+ [tx1, ty1, test1] = all_tests["non_batched_mo"]
+ model1 = SingleTaskVariationalGP(tx1, ty1, inducing_points=tx1)
+ posterior_transform = ScalarizedPosteriorTransform(
+ weights=torch.tensor([1.0, 1.0])
+ )
+ posterior1 = model1.posterior(test1, posterior_transform=posterior_transform)
+ self.assertIsInstance(posterior1, GPyTorchPosterior)
+ self.assertEqual(posterior1.mean.shape[1], 1)
+
  def test_variational_setUp(self):
  for dtype in [torch.float, torch.double]:
  train_X = torch.rand(10, 1, device=self.device, dtype=dtype)