def _transform_tensor_args(self, X: Tensor, Y: Tensor, Yvar: Optional[Tensor] = None ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: r"""Transforms tensor arguments: for single output models, the output dimension is squeezed and for multi-output models, the output dimension is transformed into the left-most batch dimension. Args: X: A `n x d` or `batch_shape x n x d` (batch mode) tensor of training features. Y: A `n x m` or `batch_shape x n x m` (batch mode) tensor of training observations. Yvar: A `n x m` or `batch_shape x n x m` (batch mode) tensor of observed measurement noise. Note: this will be None when using a model that infers the noise level (e.g. a `SingleTaskGP`). Returns: 3-element tuple containing - A `input_batch_shape x (m) x n x d` tensor of training features. - A `target_batch_shape x (m) x n` tensor of training observations. - A `target_batch_shape x (m) x n` tensor observed measurement noise (or None). """ if self._num_outputs > 1: return multioutput_to_batch_mode_transform( train_X=X, train_Y=Y, train_Yvar=Yvar, num_outputs=self._num_outputs) return X, Y.squeeze(-1), None if Yvar is None else Yvar.squeeze(-1)
def condition_on_observations( self, X: Tensor, Y: Tensor, **kwargs: Any ) -> BatchedMultiOutputGPyTorchModel: r"""Condition the model on new observations. Args: X: A `batch_shape x n' x d`-dim Tensor, where `d` is the dimension of the feature space, `m` is the number of points per batch, and `batch_shape` is the batch shape (must be compatible with the batch shape of the model). Y: A `batch_shape' x n' x m`-dim Tensor, where `m` is the number of model outputs, `n'` is the number of points per batch, and `batch_shape'` is the batch shape of the observations. `batch_shape'` must be broadcastable to `batch_shape` using standard broadcasting semantics. If `Y` has fewer batch dimensions than `X`, its is assumed that the missing batch dimensions are the same for all `Y`. Returns: A `BatchedMultiOutputGPyTorchModel` object of the same type with `n + n'` training examples, representing the original model conditioned on the new observations `(X, Y)` (and possibly noise observations passed in via kwargs). Example: >>> train_X = torch.rand(20, 2) >>> train_Y = torch.cat( >>> [torch.sin(train_X[:, 0]), torch.cos(train_X[:, 1])], -1 >>> ) >>> model = SingleTaskGP(train_X, train_Y) >>> new_X = torch.rand(5, 2) >>> new_Y = torch.cat([torch.sin(new_X[:, 0]), torch.cos(new_X[:, 1])], -1) >>> model = model.condition_on_observations(X=new_X, Y=new_Y) """ noise = kwargs.get("noise") if hasattr(self, "outcome_transform"): # we need to apply transforms before shifting batch indices around Y, noise = self.outcome_transform(Y, noise) self._validate_tensor_args(X=X, Y=Y, Yvar=noise, strict=False) inputs = X if self._num_outputs > 1: inputs, targets, noise = multioutput_to_batch_mode_transform( train_X=X, train_Y=Y, num_outputs=self._num_outputs, train_Yvar=noise ) # `multioutput_to_batch_mode_transform` removes the output dimension, # which is necessary for `condition_on_observations` targets = targets.unsqueeze(-1) if noise is not None: noise = noise.unsqueeze(-1) else: inputs = X targets = Y if noise is not None: kwargs.update({"noise": noise}) fantasy_model = super().condition_on_observations(X=inputs, Y=targets, **kwargs) fantasy_model._input_batch_shape = fantasy_model.train_targets.shape[ : (-1 if self._num_outputs == 1 else -2) ] fantasy_model._aug_batch_shape = fantasy_model.train_targets.shape[:-1] return fantasy_model
def test_multioutput_to_batch_mode_transform(self, cuda=False): for double in (False, True): tkwargs = { "device": torch.device("cuda") if cuda else torch.device("cpu"), "dtype": torch.double if double else torch.float, } n = 3 num_outputs = 1 train_X = torch.rand(n, 1, **tkwargs) train_Y = torch.rand(n, **tkwargs) train_Yvar = torch.rand(n, **tkwargs) # num_outputs = 1 and train_Y has shape `n` X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=num_outputs, train_Yvar=train_Yvar, ) self.assertTrue(torch.equal(X_out, train_X)) self.assertTrue(torch.equal(Y_out, train_Y)) self.assertTrue(torch.equal(Yvar_out, train_Yvar)) # num_outputs = 1 and train_Y has shape `n x 1` X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y.view(-1, 1), num_outputs=num_outputs, train_Yvar=train_Yvar.view(-1, 1), ) self.assertTrue(torch.equal(X_out, train_X)) self.assertTrue(torch.equal(Y_out, train_Y)) self.assertTrue(torch.equal(Yvar_out, train_Yvar)) # num_outputs > 1 num_outputs = 2 train_Y = torch.rand(n, num_outputs, **tkwargs) train_Yvar = torch.rand(n, num_outputs, **tkwargs) X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=num_outputs, train_Yvar=train_Yvar, ) expected_X_out = train_X.unsqueeze(0).expand(num_outputs, -1, 1) self.assertTrue(torch.equal(X_out, expected_X_out)) self.assertTrue(torch.equal(Y_out, train_Y.transpose(0, 1))) self.assertTrue(torch.equal(Yvar_out, train_Yvar.transpose(0, 1)))
def test_multioutput_to_batch_mode_transform(self, cuda=False): for double in (False, True): tkwargs = { "device": torch.device("cuda") if cuda else torch.device("cpu"), "dtype": torch.double if double else torch.float, } n = 3 num_outputs = 1 train_X = torch.rand(n, 1, **tkwargs) train_Y = torch.rand(n, **tkwargs) train_Yvar = torch.rand(n, **tkwargs) # num_outputs = 1 and train_Y has shape `n` X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=num_outputs, train_Yvar=train_Yvar, ) self.assertTrue(torch.equal(X_out, train_X)) self.assertTrue(torch.equal(Y_out, train_Y)) self.assertTrue(torch.equal(Yvar_out, train_Yvar)) # num_outputs = 1 and train_Y has shape `n x 1` X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y.view(-1, 1), num_outputs=num_outputs, train_Yvar=train_Yvar.view(-1, 1), ) self.assertTrue(torch.equal(X_out, train_X)) self.assertTrue(torch.equal(Y_out, train_Y)) self.assertTrue(torch.equal(Yvar_out, train_Yvar)) # num_outputs > 1 num_outputs = 2 train_Y = torch.rand(n, num_outputs, **tkwargs) train_Yvar = torch.rand(n, num_outputs, **tkwargs) X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=num_outputs, train_Yvar=train_Yvar, ) expected_X_out = train_X.unsqueeze(0).expand(num_outputs, -1, 1) self.assertTrue(torch.equal(X_out, expected_X_out)) self.assertTrue(torch.equal(Y_out, train_Y.transpose(0, 1))) self.assertTrue(torch.equal(Yvar_out, train_Yvar.transpose(0, 1)))
def __init__(self, train_X, train_Y): train_X, train_Y, _ = self._set_dimensions(train_X=train_X, train_Y=train_Y) train_X, train_Y, _ = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=self._num_outputs ) likelihood = GaussianLikelihood(batch_shape=self._aug_batch_shape) super().__init__(train_X, train_Y, likelihood) self.mean_module = ConstantMean(batch_shape=self._aug_batch_shape) self.covar_module = ScaleKernel( RBFKernel(batch_shape=self._aug_batch_shape), batch_shape=self._aug_batch_shape, )
def test_multioutput_to_batch_mode_transform(self): for dtype in (torch.float, torch.double): tkwargs = {"device": self.device, "dtype": dtype} n = 3 num_outputs = 2 train_X = torch.rand(n, 1, **tkwargs) train_Y = torch.rand(n, num_outputs, **tkwargs) train_Yvar = torch.rand(n, num_outputs, **tkwargs) X_out, Y_out, Yvar_out = multioutput_to_batch_mode_transform( train_X=train_X, train_Y=train_Y, num_outputs=num_outputs, train_Yvar=train_Yvar, ) expected_X_out = train_X.unsqueeze(0).expand(num_outputs, -1, 1) self.assertTrue(torch.equal(X_out, expected_X_out)) self.assertTrue(torch.equal(Y_out, train_Y.transpose(0, 1))) self.assertTrue(torch.equal(Yvar_out, train_Yvar.transpose(0, 1)))