Support passing a list of Tensors and floats to FixedFeatureAcquisitionFunction. (#836)

Summary:
Pull Request resolved: #836

This diff introduces an API change. Previously, one would have to instantiate a FFACQF using 1) a full Tensor or 2) a list of floats. Here, we allow passing Tensor in the list as long as they are broadcastable.

Reviewed By: Balandat

Differential Revision: D28962037

fbshipit-source-id: 526031857f969f1a9348c83a9cf1bbb47eef90f2

Author: vishwakftw

botorch/acquisition/fixed_feature.py

@@ -11,7 +11,8 @@
 
 from __future__ import annotations
 
-from typing import List, Union
+from numbers import Number
+from typing import List, Sequence, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
@@ -36,7 +37,7 @@ def __init__(
         acq_function: AcquisitionFunction,
         d: int,
         columns: List[int],
-        values: Union[Tensor, List[float]],
+        values: Union[Tensor, Sequence[Union[Tensor, float]]],
     ) -> None:
         r"""Derived Acquisition Function by fixing a subset of input features.
 
@@ -51,16 +52,47 @@ def __init__(
                 different for each of the `q` input points), or an array-like of
                 values that is broadcastable to the input across `t`-batch and
                 `q`-batch dimensions, e.g. a list of length `d_f` if values
-                are the same across all `t` and `q`-batch dimensions.
+                are the same across all `t` and `q`-batch dimensions, or a
+                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
         self.d = d
-        values = torch.as_tensor(values).detach().clone()
-        self.register_buffer("values", values)
+        if isinstance(values, Tensor):
+            new_values = values.detach().clone()
+        else:
+            new_values = []
+            for value in values:
+                if isinstance(value, Number):
+                    new_values.append(torch.tensor([float(value)]))
+                else:
+                    new_values.append(value.detach().clone())
+
+            # There are 3 cases for when `values` is a `Sequence`.
+            # 1) `values` == list of floats as earlier.
+            # 2) `values` == combination of floats and `Tensor`s.
+            # 3) `values` == a list of `Tensor`s.
+            # For 1), the below step creates a vector of length `len(values)`
+            # For 2), the below step creates a `Tensor` of shape `batch_shape x q x d_f`
+            # with the broadcasting functionality.
+            # For 3), this is simply a concatenation, yielding a `Tensor` with the
+            # same shape as in 2).
+            # The key difference arises when `_construct_X_full` is invoked.
+            # In 1), the expansion (`self.values.expand`) will expand the `Tensor` to
+            # size `batch_shape x q x d_f`.
+            # In 2) and 3), this expansion is a no-op because they are already of the
+            # required size. However, 2) and 3) _cannot_ support varying `batch_shape`,
+            # which means that all calls to `FixedFeatureAcquisitionFunction` have
+            # to have the same size throughout when `values` contains a `Tensor`.
+            # This is consistent with the scenario when a singular `Tensor` is passed
+            # as the `values` argument.
+            new_values = torch.cat(torch.broadcast_tensors(*new_values), dim=-1)
+
+        self.register_buffer("values", new_values)
         # build selector for _construct_X_full
         self._selector = []
-        idx_X, idx_f = 0, d - values.shape[-1]
+        idx_X, idx_f = 0, d - new_values.shape[-1]
         for i in range(self.d):
             if i in columns:
                 self._selector.append(idx_f)

test/acquisition/test_fixed_feature.py

@@ -17,33 +17,52 @@ def test_fixed_features(self):
         train_Y = train_X.norm(dim=-1, keepdim=True)
         model = SingleTaskGP(train_X, train_Y).to(device=self.device).eval()
         qEI = qExpectedImprovement(model, best_f=0.0)
-        # test single point
-        test_X = torch.rand(1, 3, device=self.device)
-        qEI_ff = FixedFeatureAcquisitionFunction(
-            qEI, d=3, columns=[2], values=test_X[..., -1:]
-        )
-        qei = qEI(test_X)
-        qei_ff = qEI_ff(test_X[..., :-1])
-        self.assertTrue(torch.allclose(qei, qei_ff))
-        # test list input
-        qEI_ff = FixedFeatureAcquisitionFunction(qEI, d=3, columns=[2], values=[0.5])
-        qei_ff = qEI_ff(test_X[..., :-1])
-        # test q-batch
-        test_X = torch.rand(2, 3, device=self.device)
-        qEI_ff = FixedFeatureAcquisitionFunction(
-            qEI, d=3, columns=[1], values=test_X[..., [1]]
-        )
-        qei = qEI(test_X)
-        qei_ff = qEI_ff(test_X[..., [0, 2]])
-        self.assertTrue(torch.allclose(qei, qei_ff))
-        # test t-batch with broadcasting
-        test_X = torch.rand(2, 3, device=self.device).expand(4, 2, 3)
-        qEI_ff = FixedFeatureAcquisitionFunction(
-            qEI, d=3, columns=[2], values=test_X[0, :, -1:]
-        )
-        qei = qEI(test_X)
-        qei_ff = qEI_ff(test_X[..., :-1])
-        self.assertTrue(torch.allclose(qei, qei_ff))
+        for q in [1, 2]:
+            # test single point
+            test_X = torch.rand(q, 3, device=self.device)
+            qEI_ff = FixedFeatureAcquisitionFunction(
+                qEI, d=3, columns=[2], values=test_X[..., -1:]
+            )
+            qei = qEI(test_X)
+            qei_ff = qEI_ff(test_X[..., :-1])
+            self.assertTrue(torch.allclose(qei, qei_ff))
+
+            # test list input with float
+            qEI_ff = FixedFeatureAcquisitionFunction(
+                qEI, d=3, columns=[2], values=[0.5]
+            )
+            qei_ff = qEI_ff(test_X[..., :-1])
+            test_X_clone = test_X.clone()
+            test_X_clone[..., 2] = 0.5
+            qei = qEI(test_X_clone)
+            self.assertTrue(torch.allclose(qei, qei_ff))
+
+            # test list input with Tensor and float
+            qEI_ff = FixedFeatureAcquisitionFunction(
+                qEI, d=3, columns=[0, 2], values=[test_X[..., [0]], 0.5]
+            )
+            qei_ff = qEI_ff(test_X[..., [1]])
+            self.assertTrue(torch.allclose(qei, qei_ff))
+
+            # test t-batch with broadcasting and list of floats
+            test_X = torch.rand(q, 3, device=self.device).expand(4, q, 3)
+            qEI_ff = FixedFeatureAcquisitionFunction(
+                qEI, d=3, columns=[2], values=test_X[0, :, -1:]
+            )
+            qei = qEI(test_X)
+            qei_ff = qEI_ff(test_X[..., :-1])
+            self.assertTrue(torch.allclose(qei, qei_ff))
+
+            # test t-batch with broadcasting and list of floats and Tensor
+            qEI_ff = FixedFeatureAcquisitionFunction(
+                qEI, d=3, columns=[0, 2], values=[test_X[0, :, [0]], 0.5]
+            )
+            test_X_clone = test_X.clone()
+            test_X_clone[..., 2] = 0.5
+            qei = qEI(test_X_clone)
+            qei_ff = qEI_ff(test_X[..., [1]])
+            self.assertTrue(torch.allclose(qei, qei_ff))
+
         # test gradient
         test_X = torch.rand(1, 3, device=self.device, requires_grad=True)
         test_X_ff = test_X[..., :-1].detach().clone().requires_grad_(True)
@@ -56,6 +75,20 @@ def test_fixed_features(self):
         qei.backward()
         qei_ff.backward()
         self.assertTrue(torch.allclose(test_X.grad[..., :-1], test_X_ff.grad))
+
+        test_X = test_X.detach().clone()
+        test_X_ff = test_X[..., [1]].detach().clone().requires_grad_(True)
+        test_X[..., 2] = 0.5
+        test_X.requires_grad_(True)
+        qei = qEI(test_X)
+        qEI_ff = FixedFeatureAcquisitionFunction(
+            qEI, d=3, columns=[0, 2], values=[test_X[..., [0]].detach(), 0.5]
+        )
+        qei_ff = qEI_ff(test_X_ff)
+        qei.backward()
+        qei_ff.backward()
+        self.assertTrue(torch.allclose(test_X.grad[..., [1]], test_X_ff.grad))
+
         # test error b/c of incompatible input shapes
         with self.assertRaises(ValueError):
             qEI_ff(test_X)