Add is_one_to_many attribute to input transforms (#1396)

Summary:
Pull Request resolved: #1396

Adds a simple attribute to replace `isinstance(tf, (AppendFeatures, InputPerturbation))` checks. For chained transforms, this is set to true if at least one of the child transforms is one-to-many.

Reviewed By: Balandat

Differential Revision: D39455329

fbshipit-source-id: 3381b6299c3f83a4b68fbf0d579481ea2ffd0417

Author: saitcakmak

botorch/models/transforms/input.py

@@ -41,6 +41,8 @@ class InputTransform(ABC):
         between `gpytorch.module.Module` and `torch.nn.Module` in `Warp`.
 
     Properties:
+        is_one_to_many: A boolean denoting whether the transform produces
+            multiple values for each input.
         transform_on_train: A boolean indicating whether to apply the
             transform in train() mode.
         transform_on_eval: A boolean indicating whether to apply the
@@ -51,6 +53,7 @@ class InputTransform(ABC):
     :meta private:
     """
 
+    is_one_to_many: bool = False
     transform_on_eval: bool
     transform_on_train: bool
     transform_on_fantasize: bool
@@ -177,6 +180,7 @@ def __init__(self, **transforms: InputTransform) -> None:
         self.transform_on_eval = False
         self.transform_on_fantasize = False
         for tf in transforms.values():
+            self.is_one_to_many |= tf.is_one_to_many
             self.transform_on_train |= tf.transform_on_train
             self.transform_on_eval |= tf.transform_on_eval
             self.transform_on_fantasize |= tf.transform_on_fantasize
@@ -999,6 +1003,8 @@ class AppendFeatures(InputTransform, Module):
         >>> risk_measure_samples = risk_measure(posterior_samples)
     """
 
+    is_one_to_many: bool = True
+
     def __init__(
         self,
         feature_set: Optional[Tensor] = None,
@@ -1196,6 +1202,8 @@ class InputPerturbation(InputTransform, Module):
     https://botorch.org/tutorials/risk_averse_bo_with_input_perturbations.
     """
 
+    is_one_to_many: bool = True
+
     def __init__(
         self,
         perturbation_set: Union[Tensor, Callable[[Tensor], Tensor]],

test/models/transforms/test_input.py

@@ -126,11 +126,10 @@ def test_abstract_base_input_transform(self):
             self.assertTrue(torch.equal(ipt5(X), X))
 
     def test_normalize(self):
-
         for dtype in (torch.float, torch.double):
-
             # basic init, learned bounds
             nlz = Normalize(d=2)
+            self.assertFalse(nlz.is_one_to_many)
             self.assertTrue(nlz.learn_bounds)
             self.assertTrue(nlz.training)
             self.assertEqual(nlz._d, 2)
@@ -316,7 +315,6 @@ def test_normalize(self):
 
     def test_standardize(self):
         for dtype in (torch.float, torch.double):
-
             # basic init
             stdz = InputStandardize(d=2)
             self.assertTrue(stdz.training)
@@ -455,7 +453,6 @@ def test_standardize(self):
                 self.assertFalse(stdz7.equals(stdz8))
 
     def test_chained_input_transform(self):
-
         ds = (1, 2)
         batch_shapes = (torch.Size(), torch.Size([2]))
         dtypes = (torch.float, torch.double)
@@ -473,6 +470,7 @@ def test_chained_input_transform(self):
             self.assertEqual(sorted(tf.keys()), ["stz_fixed", "stz_learned"])
             self.assertEqual(tf["stz_fixed"], tf1)
             self.assertEqual(tf["stz_learned"], tf2)
+            self.assertFalse(tf.is_one_to_many)
 
             X = torch.rand(*batch_shape, 4, d, device=self.device, dtype=dtype)
             X_tf = tf(X)
@@ -515,6 +513,11 @@ def test_chained_input_transform(self):
             tf = ChainedInputTransform(stz_fixed=tf1, stz_learned=tf2)
             self.assertTrue(torch.equal(tf.preprocess_transform(X), tf1.transform(X)))
 
+        # test one-to-many
+        tf2 = InputPerturbation(perturbation_set=bounds)
+        tf = ChainedInputTransform(stz=tf1, pert=tf2)
+        self.assertTrue(tf.is_one_to_many)
+
     def test_round_transform(self):
         for dtype in (torch.float, torch.double):
             # basic init
@@ -811,6 +814,7 @@ def test_append_features(self):
                 torch.linspace(0, 1, 6).view(3, 2).to(device=self.device, dtype=dtype)
             )
             transform = AppendFeatures(feature_set=feature_set)
+            self.assertTrue(transform.is_one_to_many)
             X = torch.rand(4, 5, 3, device=self.device, dtype=dtype)
             # in train - no transform
             transform.train()
@@ -1173,6 +1177,7 @@ def test_input_perturbation(self):
                 [[0.5, -0.3], [0.2, 0.4], [-0.7, 0.1]], device=self.device, dtype=dtype
             )
             transform = InputPerturbation(perturbation_set=perturbation_set)
+            self.assertTrue(transform.is_one_to_many)
             X = torch.tensor(
                 [[[0.5, 0.5], [0.9, 0.7]], [[0.3, 0.2], [0.1, 0.4]]],
                 device=self.device,