fix input constructor for LogProbabilityOfFeasibility

Summary:
X-link: facebook/Ax#4080

Constraints were being passed in the wrong format (a list of callables rather than a dictionary of indices to bounds).

This also removes `convert_to_target_pre_hook`, which only works for args and doesn't work for kwargs. Previously `LogProbabilityOfFeasibility` would error out if called with a kwarg---e.g. `acqf(X=X)` dtype and device are set in `_mean_and_sigma` for these AFs anyway.

Differential Revision: D79281274

Author: sdaulton

botorch/acquisition/analytic.py

@@ -37,7 +37,6 @@
 from botorch.utils.safe_math import log1mexp, logmeanexp
 from botorch.utils.transforms import (
     average_over_ensemble_models,
-    convert_to_target_pre_hook,
     t_batch_mode_transform,
 )
 from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
@@ -100,7 +99,7 @@ def _mean_and_sigma(
             posterior. Removes the last two dimensions if they have size one. Only
             returns a single tensor of means if compute_sigma is True.
         """
-        self.to(device=X.device)  # ensures buffers / parameters are on the same device
+        self.to(X)  # ensures buffers / parameters are on the same device and dtype
         posterior = self.model.posterior(
             X=X, posterior_transform=self.posterior_transform
         )
@@ -584,7 +583,6 @@ def __init__(
         self.objective_index = objective_index
         self.register_buffer("best_f", torch.as_tensor(best_f))
         ConstrainedAnalyticAcquisitionFunctionMixin.__init__(self, constraints)
-        self.register_forward_pre_hook(convert_to_target_pre_hook)
 
     @t_batch_mode_transform(expected_q=1)
     @average_over_ensemble_models
@@ -638,9 +636,7 @@ class LogProbabilityOfFeasibility(
     _log: bool = True
 
     def __init__(
-        self,
-        model: Model,
-        constraints: dict[int, tuple[float | None, float | None]],
+        self, model: Model, constraints: dict[int, tuple[float | None, float | None]]
     ) -> None:
         r"""Analytic Log Probability of Feasibility.
 
@@ -654,7 +650,6 @@ def __init__(
         AcquisitionFunction.__init__(self, model=model)
         self.posterior_transform = None
         ConstrainedAnalyticAcquisitionFunctionMixin.__init__(self, constraints)
-        self.register_forward_pre_hook(convert_to_target_pre_hook)
 
     @t_batch_mode_transform(expected_q=1)
     @average_over_ensemble_models
@@ -730,7 +725,6 @@ def __init__(
         self.objective_index = objective_index
         self.register_buffer("best_f", torch.as_tensor(best_f))
         ConstrainedAnalyticAcquisitionFunctionMixin.__init__(self, constraints)
-        self.register_forward_pre_hook(convert_to_target_pre_hook)
 
     @t_batch_mode_transform(expected_q=1)
     @average_over_ensemble_models

botorch/acquisition/input_constructors.py

@@ -97,7 +97,7 @@
     get_optimal_samples,
     project_to_target_fidelity,
 )
-from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.errors import BotorchError, UnsupportedError
 from botorch.models.cost import AffineFidelityCostModel
 from botorch.models.deterministic import FixedSingleSampleModel
 from botorch.models.gpytorch import GPyTorchModel
@@ -226,6 +226,12 @@ def allow_only_specific_variable_kwargs(f: Callable[..., T]) -> Callable[..., T]
         # Used in input constructors for some lookahead acquisition functions
         # such as qKnowledgeGradient.
         "bounds",
+        # Needed for LogProbabilityOfFeasibility
+        "constraints_tuple",
+        "posterior_transform",
+        # not used by analytic acquisition functions
+        "objective",
+        "constraints",
     }
 
     def g(*args: Any, **kwargs: Any) -> T:
@@ -338,28 +344,42 @@ def construct_inputs_best_f(
     }
 
 
-@acqf_input_constructor(
-    LogProbabilityOfFeasibility,
-)
+@acqf_input_constructor(LogProbabilityOfFeasibility)
 def construct_inputs_pof(
-    model: Model,
-    constraints: dict[int, tuple[float | None, float | None]],
+    model: Model, constraints_tuple: tuple[Tensor, Tensor]
 ) -> dict[str, Any]:
     r"""Construct kwargs for the log probability of feasibility acquisition function.
 
     Args:
         model: The model to be used in the acquisition function.
-        constraints: A dictionary of the form `{i: [lower, upper]}`, where `i` is the
-            output index, and `lower` and `upper` are lower and upper bounds on that
-            output (resp. interpreted as -Inf / Inf if None).
+        constraints_tuple: A tuple of `(A, b)`. For `k` outcome constraints
+            and `m` outputs at `f(x)``, `A` is `k x m` and `b` is `k x 1` such
+            that `A f(x) <= b`.
+
 
     Returns:
         A dict mapping kwarg names of the constructor to values.
     """
-    return {
-        "model": model,
-        "constraints": constraints,
-    }
+    # Construct a dictionary of the form `{i: [lower, upper]}`,
+    # where `i` is the output index, and `lower` and `upper` are
+    # lower and upper bounds on that output (resp. interpreted
+    # as -Inf / Inf if None).
+    weights, bounds = constraints_tuple
+    constraints_dict = {}
+    for w, b in zip(weights, bounds):
+        nonzero_w = w.nonzero()
+        if nonzero_w.numel() != 1:
+            raise BotorchError(
+                "LogProbabilityOfFeasibility only support constraints on single"
+                " outcomes."
+            )
+        i = nonzero_w.item()
+        w_i = w[i]
+        is_ub = torch.sign(w_i) == 1.0
+        b = b.item()
+        bounds = (None, b / w_i) if is_ub else (b / w_i, None)
+        constraints_dict[i] = bounds
+    return {"model": model, "constraints": constraints_dict}
 
 
 @acqf_input_constructor(UpperConfidenceBound)

botorch/utils/transforms.py

@@ -422,8 +422,3 @@ def match_batch_shape(X: Tensor, Y: Tensor) -> Tensor:
 
     """
     return X.expand(X.shape[: -(Y.dim())] + Y.shape[:-2] + X.shape[-2:])
-
-
-def convert_to_target_pre_hook(module, *args):
-    r"""Pre-hook for automatically calling `.to(X)` on module prior to `forward`"""
-    module.to(args[0][0])

test/acquisition/test_input_constructors.py

@@ -110,7 +110,7 @@
     expand_trace_observations,
     project_to_target_fidelity,
 )
-from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.errors import BotorchError, UnsupportedError
 from botorch.models import MultiTaskGP, SaasFullyBayesianSingleTaskGP, SingleTaskGP
 from botorch.models.deterministic import FixedSingleSampleModel
 from botorch.models.model_list_gp_regression import ModelListGP
@@ -636,11 +636,26 @@ def constraint(Y: Tensor) -> Tensor:
     def test_construct_inputs_LogPOF(self) -> None:
         c = get_acqf_input_constructor(LogProbabilityOfFeasibility)
         mock_model = self.mock_model
-        constraints = {1: [None, 0]}
-        kwargs = c(model=mock_model, constraints=constraints)
+        constraints_tuple = [torch.tensor([[0.0, 1.0]]), torch.tensor([[2.0]])]
+        constraints = {1: (None, 2.0)}
+        kwargs = c(model=mock_model, constraints_tuple=constraints_tuple)
         self.assertEqual(set(kwargs.keys()), {"model", "constraints"})
         self.assertIs(kwargs["model"], mock_model)
         self.assertEqual(kwargs["constraints"], constraints)
+        constraints_tuple = [torch.tensor([[0.0, -1.0]]), torch.tensor([[-2.0]])]
+        kwargs = c(model=mock_model, constraints_tuple=constraints_tuple)
+        constraints = {1: (2.0, None)}
+        self.assertEqual(kwargs["constraints"], constraints)
+        # test that constraints on multiple outcomes raises an exception
+        with self.assertRaisesRegex(
+            BotorchError,
+            "LogProbabilityOfFeasibility only support constraints on single"
+            " outcomes.",
+        ):
+            c(
+                model=mock_model,
+                constraints_tuple=[torch.tensor([[1.0, 1.0]]), torch.tensor([[2.0]])],
+            )
 
     def test_construct_inputs_qEI(self) -> None:
         c = get_acqf_input_constructor(qExpectedImprovement)