Probabilistic Reparameterization Part 2

botorch/acquisition/__init__.py

@@ -80,12 +80,17 @@
     qExpectedUtilityOfBestOption,
 )
 from botorch.acquisition.prior_guided import PriorGuidedAcquisitionFunction
+from botorch.acquisition.probabilistic_reparameterization import (
+    AnalyticProbabilisticReparameterization,
+    MCProbabilisticReparameterization,
+)
 from botorch.acquisition.proximal import ProximalAcquisitionFunction
 
 __all__ = [
     "AcquisitionFunction",
     "AnalyticAcquisitionFunction",
     "AnalyticExpectedUtilityOfBestOption",
+    "AnalyticProbabilisticReparameterization",
     "ConstrainedExpectedImprovement",
     "DecoupledAcquisitionFunction",
     "ExpectedImprovement",
@@ -94,6 +99,7 @@
     "FixedFeatureAcquisitionFunction",
     "GenericCostAwareUtility",
     "InverseCostWeightedUtility",
+    "MCProbabilisticReparameterization",
     "NoisyExpectedImprovement",
     "OneShotAcquisitionFunction",
     "PairwiseBayesianActiveLearningByDisagreement",

botorch/acquisition/fixed_feature.py

@@ -17,8 +17,8 @@
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
+from botorch.acquisition.wrapper import AbstractAcquisitionFunctionWrapper
 from torch import Tensor
-from torch.nn import Module
 
 
 def get_dtype_of_sequence(values: Sequence[Tensor | float]) -> torch.dtype:
@@ -50,8 +50,8 @@ def _is_cuda(value: Tensor | float) -> bool:
     return torch.device("cuda") if any_cuda else torch.device("cpu")
 
 
-class FixedFeatureAcquisitionFunction(AcquisitionFunction):
-    """A wrapper around AcquisitionFunctions to fix a subset of features.
+class FixedFeatureAcquisitionFunction(AbstractAcquisitionFunctionWrapper):
+    """A wrapper around AquisitionFunctions to fix a subset of features.
 
     Example:
         >>> model = SingleTaskGP(train_X, train_Y)  # d = 5
@@ -86,8 +86,9 @@ def __init__(
                 combination of `Tensor`s and numbers which can be broadcasted
                 to form a tensor with trailing dimension size of `d_f`.
         """
-        Module.__init__(self)
-        self.acq_func = acq_function
+        AbstractAcquisitionFunctionWrapper.__init__(self, acq_function=acq_function)
+        dtype = torch.float
+        device = torch.device("cpu")
         self.d = d
 
         if isinstance(values, Tensor):
@@ -153,24 +154,13 @@ def forward(self, X: Tensor):
         X_full = self._construct_X_full(X)
         return self.acq_func(X_full)
 
-    @property
-    def X_pending(self):
-        r"""Return the `X_pending` of the base acquisition function."""
-        try:
-            return self.acq_func.X_pending
-        except (ValueError, AttributeError):
-            raise ValueError(
-                f"Base acquisition function {type(self.acq_func).__name__} "
-                "does not have an `X_pending` attribute."
-            )
-
-    @X_pending.setter
-    def X_pending(self, X_pending: Tensor | None):
+    def set_X_pending(self, X_pending: Tensor | None):
         r"""Sets the `X_pending` of the base acquisition function."""
         if X_pending is not None:
-            self.acq_func.X_pending = self._construct_X_full(X_pending)
+            full_X_pending = self._construct_X_full(X_pending)
         else:
-            self.acq_func.X_pending = X_pending
+            full_X_pending = None
+        self.acq_func.set_X_pending(full_X_pending)
 
     def _construct_X_full(self, X: Tensor) -> Tensor:
         r"""Constructs the full input for the base acquisition function.

botorch/acquisition/penalized.py

@@ -16,9 +16,8 @@
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
-from botorch.acquisition.analytic import AnalyticAcquisitionFunction
 from botorch.acquisition.objective import GenericMCObjective
-from botorch.exceptions import UnsupportedError
+from botorch.acquisition.wrapper import AbstractAcquisitionFunctionWrapper
 from torch import Tensor
 
 
@@ -201,7 +200,7 @@ def __call__(self, X: Tensor) -> Tensor:
         return super().__call__(X=X).squeeze(dim=-1).min(dim=-1).values
 
 
-class PenalizedAcquisitionFunction(AcquisitionFunction):
+class PenalizedAcquisitionFunction(AbstractAcquisitionFunctionWrapper):
     r"""Single-outcome acquisition function regularized by the given penalty.
 
     The usage is similar to:
@@ -223,29 +222,16 @@ def __init__(
             penalty_func: The regularization function.
             regularization_parameter: Regularization parameter used in optimization.
         """
-        super().__init__(model=raw_acqf.model)
-        self.raw_acqf = raw_acqf
+        AcquisitionFunction.__init__(self, model=raw_acqf.model)
+        AbstractAcquisitionFunctionWrapper.__init__(self, acq_function=raw_acqf)
         self.penalty_func = penalty_func
         self.regularization_parameter = regularization_parameter
 
     def forward(self, X: Tensor) -> Tensor:
-        raw_value = self.raw_acqf(X=X)
+        raw_value = self.acq_func(X=X)
         penalty_term = self.penalty_func(X)
         return raw_value - self.regularization_parameter * penalty_term
 
-    @property
-    def X_pending(self) -> Tensor | None:
-        return self.raw_acqf.X_pending
-
-    def set_X_pending(self, X_pending: Tensor | None = None) -> None:
-        if not isinstance(self.raw_acqf, AnalyticAcquisitionFunction):
-            self.raw_acqf.set_X_pending(X_pending=X_pending)
-        else:
-            raise UnsupportedError(
-                "The raw acquisition function is Analytic and does not account "
-                "for X_pending yet."
-            )
-
 
 def group_lasso_regularizer(X: Tensor, groups: list[list[int]]) -> Tensor:
     r"""Computes the group lasso regularization function for the given point.