Add a skip_expand option to AppendFeatures (#1344)

Summary:
Pull Request resolved: #1344

Adds a `skip_expand` option to `AppendFeatures` transform, which skips the `expand` operation in the `forward` pass, allowing it to be mixed with an `InputPerturbation` transform without multiplying the size of the output tensor. When used together, each perturbation and each appended feature will appear only once in the output tensor per q-batch.

Reviewed By: danielrjiang

Differential Revision: D38594609

fbshipit-source-id: 8cf1a8cbfb1be6faf3ba7a524d784f6488f461f0

Author: saitcakmak

botorch/models/transforms/input.py

@@ -965,6 +965,7 @@ class AppendFeatures(InputTransform, Module):
     def __init__(
         self,
         feature_set: Tensor,
+        skip_expand: bool = False,
         transform_on_train: bool = False,
         transform_on_eval: bool = True,
         transform_on_fantasize: bool = False,
@@ -974,6 +975,10 @@ def __init__(
         Args:
             feature_set: An `n_f x d_f`-dim tensor denoting the features to be
                 appended to the inputs.
+            skip_expand: A boolean indicating whether to expand the input tensor
+                before appending features. This is intended for use with an
+                `InputPerturbation`. If `True`, the input tensor will be expected
+                to be of shape `batch_shape x (q * n_f) x d`.
             transform_on_train: A boolean indicating whether to apply the
                 transforms in train() mode. Default: False.
             transform_on_eval: A boolean indicating whether to apply the
@@ -984,6 +989,7 @@ def __init__(
         super().__init__()
         if feature_set.dim() != 2:
             raise ValueError("`feature_set` must be an `n_f x d_f`-dim tensor!")
+        self.skip_expand = skip_expand
         self.register_buffer("feature_set", feature_set)
         self.transform_on_train = transform_on_train
         self.transform_on_eval = transform_on_eval
@@ -1003,14 +1009,22 @@ def transform(self, X: Tensor) -> Tensor:
         sample paths.
 
         Args:
-            X: A `batch_shape x q x d`-dim tensor of inputs.
+            X: A `batch_shape x q x d`-dim tensor of inputs. If `self.skip_expand` is
+                `True`, then `X` should be of shape `batch_shape x (q * n_f) x d`,
+                typically obtained by passing a `batch_shape x q x d` shape input
+                through an `InputPerturbation` with `n_f` perturbation values.
 
         Returns:
             A `batch_shape x (q * n_f) x (d + d_f)`-dim tensor of appended inputs.
         """
-        expanded_X = X.unsqueeze(dim=-2).expand(
-            *X.shape[:-1], self.feature_set.shape[0], -1
-        )
+        if self.skip_expand:
+            expanded_X = X.view(
+                *X.shape[:-2], -1, self.feature_set.shape[0], X.shape[-1]
+            )
+        else:
+            expanded_X = X.unsqueeze(dim=-2).expand(
+                *X.shape[:-1], self.feature_set.shape[0], -1
+            )
         expanded_features = self.feature_set.expand(*expanded_X.shape[:-1], -1)
         appended_X = torch.cat([expanded_X, expanded_features], dim=-1)
         return appended_X.view(*X.shape[:-2], -1, appended_X.shape[-1])

test/models/transforms/test_input.py

@@ -823,6 +823,29 @@ def test_append_features(self):
             self.assertEqual(transform.feature_set.device.type, "cpu")
             self.assertEqual(transform.feature_set.dtype, torch.half)
 
+    def test_w_skip_expand(self):
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            feature_set = torch.tensor([[0.0], [1.0]], **tkwargs)
+            append_tf = AppendFeatures(feature_set=feature_set, skip_expand=True).eval()
+            perturbation_set = torch.tensor([[0.0, 0.5], [1.0, 1.5]], **tkwargs)
+            pert_tf = InputPerturbation(perturbation_set=perturbation_set).eval()
+            test_X = torch.tensor([[0.0, 0.0], [1.0, 1.0]], **tkwargs)
+            tf_X = append_tf(pert_tf(test_X))
+            expected_X = torch.tensor(
+                [
+                    [0.0, 0.5, 0.0],
+                    [1.0, 1.5, 1.0],
+                    [1.0, 1.5, 0.0],
+                    [2.0, 2.5, 1.0],
+                ],
+                **tkwargs,
+            )
+            self.assertTrue(torch.allclose(tf_X, expected_X))
+            # Batched evaluation.
+            tf_X = append_tf(pert_tf(test_X.expand(3, 5, -1, -1)))
+            self.assertTrue(torch.allclose(tf_X, expected_X.expand(3, 5, -1, -1)))
+
 
 class TestFilterFeatures(BotorchTestCase):
     def test_filter_features(self):