Add batch_shape property to GP model class

Implements #2301.

TODO:
- support for approximate GP models
- unit tests
- verify compatibility with the botorch setup of other models

Author: Balandat

gpytorch/models/exact_gp.py

@@ -71,6 +71,17 @@ def __init__(self, train_inputs, train_targets, likelihood):
 
         self.prediction_strategy = None
 
+    @property
+    def batch_shape(self) -> torch.Size:
+        r"""The batch shape of the model.
+
+        This is a batch shape from an I/O perspective, independent of the internal
+        representation of the model. For a model with `(m)` outputs, a
+        `test_batch_shape x q x d`-shaped input to the model in eval mode returns a
+        distribution of shape `broadcast(test_batch_shape, model.batch_shape) x q x (m)`.
+        """
+        return self.train_inputs[0].shape[:-2]
+
     @property
     def train_targets(self):
         return self._train_targets
@@ -160,9 +171,7 @@ def get_fantasy_model(self, inputs, targets, **kwargs):
                 "all test independent caches exist. Call the model on some data first!"
             )
 
-        model_batch_shape = self.train_inputs[0].shape[:-2]
-
-        if self.train_targets.dim() > len(model_batch_shape) + 1:
+        if self.train_targets.dim() > len(self.batch_shape) + 1:
             raise RuntimeError("Cannot yet add fantasy observations to multitask GPs, but this is coming soon!")
 
         if not isinstance(inputs, list):
@@ -182,17 +191,17 @@ def get_fantasy_model(self, inputs, targets, **kwargs):
 
         # Check whether we can properly broadcast batch dimensions
         try:
-            torch.broadcast_shapes(model_batch_shape, target_batch_shape)
+            torch.broadcast_shapes(self.batch_shape, target_batch_shape)
         except RuntimeError:
             raise RuntimeError(
-                f"Model batch shape ({model_batch_shape}) and target batch shape "
+                f"Model batch shape ({self.batch_shape}) and target batch shape "
                 f"({target_batch_shape}) are not broadcastable."
             )
 
-        if len(model_batch_shape) > len(input_batch_shape):
-            input_batch_shape = model_batch_shape
-        if len(model_batch_shape) > len(target_batch_shape):
-            target_batch_shape = model_batch_shape
+        if len(self.batch_shape) > len(input_batch_shape):
+            input_batch_shape = self.batch_shape
+        if len(self.batch_shape) > len(target_batch_shape):
+            target_batch_shape = self.batch_shape
 
         # If input has no fantasy batch dimension but target does, we can save memory and computation by not
         # computing the covariance for each element of the batch. Therefore we don't expand the inputs to the

gpytorch/models/gp.py

@@ -1,7 +1,19 @@
 #!/usr/bin/env python3
 
+import torch
+
 from ..module import Module
 
 
 class GP(Module):
-    pass
+    @property
+    def batch_shape(self) -> torch.Size:
+        r"""The batch shape of the model.
+
+        This is a batch shape from an I/O perspective, independent of the internal
+        representation of the model. For a model with `(m)` outputs, a
+        `test_batch_shape x q x d`-shaped input to the model in eval mode returns a
+        distribution of shape `broadcast(test_batch_shape, model.batch_shape) x q x (m)`.
+        """
+        cls_name = self.__class__.__name__
+        raise NotImplementedError(f"{cls_name} does not define batch_shape property")

gpytorch/models/model_list.py

@@ -40,6 +40,24 @@ def __init__(self, *models):
     def num_outputs(self):
         return len(self.models)
 
+    @property
+    def batch_shape(self) -> torch.Size:
+        r"""The batch shape of the model.
+
+        This is a batch shape from an I/O perspective, independent of the internal
+        representation of the model. For a model with `(m)` outputs, a
+        `test_batch_shape x q x d`-shaped input to the model in eval mode returns a
+        distribution of shape `broadcast(test_batch_shape, model.batch_shape) x q x (m)`.
+        """
+        batch_shape = self.models[0].batch_shape
+        if all(batch_shape == m.batch_shape for m in self.models[1:]):
+            return batch_shape
+        # TODO: Allow broadcasting of model batch shapes
+        raise NotImplementedError(
+            f"`{self.__class__.__name__}.batch_shape` is only supported if all "
+            "constituent models have the same `batch_shape`."
+        )
+
     def forward_i(self, i, *args, **kwargs):
         return self.models[i].forward(*args, **kwargs)