static method for getting batch shapes for batched MO models (#346)

Summary:
Pull Request resolved: #346

This method makes it easy to figure out the `aug_batch_shape`, which is useful when specifying custom `covar_module`s and `likelihood`s.

Reviewed By: Balandat

Differential Revision: D19252473

fbshipit-source-id: 7ab7bb1a77016840fa14994810c09fa03f1821bc

Author: sdaulton

botorch/models/gpytorch.py

@@ -194,6 +194,31 @@ class BatchedMultiOutputGPyTorchModel(GPyTorchModel):
     _input_batch_shape: torch.Size
     _aug_batch_shape: torch.Size
 
+    @staticmethod
+    def get_batch_dimensions(
+        train_X: Tensor, train_Y: Tensor
+    ) -> Tuple[torch.Size, torch.Size]:
+        r"""Get the raw batch shape and output-augmented batch shape of the inputs.
+
+        Args:
+            train_X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of training
+                features.
+            train_Y: A `n x m` or `batch_shape x n x m` (batch mode) tensor of
+                training observations.
+
+        Returns:
+            2-element tuple containing
+
+            - The `input_batch_shape`
+            - The output-augmented batch shape: `input_batch_shape x (m)`
+        """
+        input_batch_shape = train_X.shape[:-2]
+        aug_batch_shape = input_batch_shape
+        num_outputs = train_Y.shape[-1]
+        if num_outputs > 1:
+            aug_batch_shape += torch.Size([num_outputs])
+        return input_batch_shape, aug_batch_shape
+
     def _set_dimensions(self, train_X: Tensor, train_Y: Tensor) -> None:
         r"""Store the number of outputs and the batch shape.
 
@@ -204,10 +229,9 @@ def _set_dimensions(self, train_X: Tensor, train_Y: Tensor) -> None:
                 training observations.
         """
         self._num_outputs = train_Y.shape[-1]
-        self._input_batch_shape = train_X.shape[:-2]
-        self._aug_batch_shape = self._input_batch_shape
-        if self._num_outputs > 1:
-            self._aug_batch_shape += torch.Size([self._num_outputs])
+        self._input_batch_shape, self._aug_batch_shape = self.get_batch_dimensions(
+            train_X=train_X, train_Y=train_Y
+        )
 
     def _transform_tensor_args(
         self, X: Tensor, Y: Tensor, Yvar: Optional[Tensor] = None

test/models/test_gpytorch.py

@@ -221,6 +221,27 @@ def test_batched_multi_output_gpytorch_model(self):
             self.assertIsInstance(cm, SimpleBatchedMultiOutputGPyTorchModel)
             self.assertEqual(cm.train_targets.shape, torch.Size([2, 2, 7]))
 
+            # test get_batch_dimensions
+            get_batch_dims = SimpleBatchedMultiOutputGPyTorchModel.get_batch_dimensions
+            for input_batch_dim in (0, 3):
+                for num_outputs in (1, 2):
+                    input_batch_shape, aug_batch_shape = get_batch_dims(
+                        train_X=train_X.unsqueeze(0).expand(3, 5, 1)
+                        if input_batch_dim == 3
+                        else train_X,
+                        train_Y=train_Y[:, 0:1] if num_outputs == 1 else train_Y,
+                    )
+                    expected_input_batch_shape = (
+                        torch.Size([3]) if input_batch_dim == 3 else torch.Size([])
+                    )
+                    self.assertEqual(input_batch_shape, expected_input_batch_shape)
+                    self.assertEqual(
+                        aug_batch_shape,
+                        expected_input_batch_shape + torch.Size([])
+                        if num_outputs == 1
+                        else expected_input_batch_shape + torch.Size([2]),
+                    )
+
 
 class TestModelListGPyTorchModel(BotorchTestCase):
     def test_model_list_gpytorch_model(self):