def test_add_output_dim(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, } original_batch_shape = torch.Size([2]) # check exception is raised X = torch.rand(2, 1, **tkwargs) with self.assertRaises(ValueError): add_output_dim(X=X, original_batch_shape=original_batch_shape) # test no new batch dims X = torch.rand(2, 2, 1, **tkwargs) X_out, output_dim_idx = add_output_dim( X=X, original_batch_shape=original_batch_shape ) self.assertTrue(torch.equal(X_out, X.unsqueeze(0))) self.assertEqual(output_dim_idx, 0) # test new batch dims X = torch.rand(3, 2, 2, 1, **tkwargs) X_out, output_dim_idx = add_output_dim( X=X, original_batch_shape=original_batch_shape ) self.assertTrue(torch.equal(X_out, X.unsqueeze(1))) self.assertEqual(output_dim_idx, 1)
def _get_pvar_expected(posterior, model, X, m): lh_kwargs = {} if isinstance(model.likelihood, FixedNoiseGaussianLikelihood): lh_kwargs["noise"] = model.likelihood.noise.mean().expand(X.shape[:-1]) if m == 1: return model.likelihood(posterior.mvn, X, **lh_kwargs).variance.unsqueeze(-1) X_, odi = add_output_dim(X=X, original_batch_shape=model._input_batch_shape) pvar_exp = model.likelihood(model(X_), X_, **lh_kwargs).variance return torch.stack([pvar_exp.select(dim=odi, index=i) for i in range(m)], dim=-1)
def _get_pvar_expected(posterior, model, X, num_outputs): if num_outputs == 1: return model.likelihood(posterior.mvn, X).variance.unsqueeze(-1) X_, odi = add_output_dim(X=X, original_batch_shape=model._input_batch_shape) pvar_exp = model.likelihood(model(X_), X_).variance return torch.stack( [pvar_exp.select(dim=odi, index=i) for i in range(num_outputs)], dim=-1)
def test_add_output_dim(self): for dtype in (torch.float, torch.double): tkwargs = {"device": self.device, "dtype": dtype} original_batch_shape = torch.Size([2]) # check exception is raised when trailing batch dims do not line up X = torch.rand(2, 3, 2, 1, **tkwargs) with self.assertRaises(RuntimeError): add_output_dim(X=X, original_batch_shape=original_batch_shape) # test no new batch dims X = torch.rand(2, 2, 1, **tkwargs) X_out, output_dim_idx = add_output_dim( X=X, original_batch_shape=original_batch_shape) self.assertTrue(torch.equal(X_out, X.unsqueeze(1))) self.assertEqual(output_dim_idx, 1) # test new batch dims X = torch.rand(3, 2, 2, 1, **tkwargs) X_out, output_dim_idx = add_output_dim( X=X, original_batch_shape=original_batch_shape) self.assertTrue(torch.equal(X_out, X.unsqueeze(2))) self.assertEqual(output_dim_idx, 2)
def posterior( self, X: Tensor, output_indices: Optional[List[int]] = None, observation_noise: Union[bool, Tensor] = False, **kwargs: Any, ) -> GPyTorchPosterior: r"""Computes the posterior over model outputs at the provided points. Args: X: A `(batch_shape) x q x d`-dim Tensor, where `d` is the dimension of the feature space and `q` is the number of points considered jointly. output_indices: A list of indices, corresponding to the outputs over which to compute the posterior (if the model is multi-output). Can be used to speed up computation if only a subset of the model's outputs are required for optimization. If omitted, computes the posterior over all model outputs. observation_noise: If True, add the observation noise from the likelihood to the posterior. If a Tensor, use it directly as the observation noise (must be of shape `(batch_shape) x q x m`). Returns: A `GPyTorchPosterior` object, representing `batch_shape` joint distributions over `q` points and the outputs selected by `output_indices` each. Includes observation noise if specified. """ self.eval() # make sure model is in eval mode with gpt_posterior_settings(): # insert a dimension for the output dimension if self._num_outputs > 1: X, output_dim_idx = add_output_dim( X=X, original_batch_shape=self._input_batch_shape) mvn = self(X) if observation_noise is not False: if torch.is_tensor(observation_noise): # TODO: Validate noise shape # make observation_noise `batch_shape x q x n` obs_noise = observation_noise.transpose(-1, -2) mvn = self.likelihood(mvn, X, noise=obs_noise) elif isinstance(self.likelihood, FixedNoiseGaussianLikelihood): # Use the mean of the previous noise values (TODO: be smarter here). noise = self.likelihood.noise.mean().expand(X.shape[:-1]) mvn = self.likelihood(mvn, X, noise=noise) else: mvn = self.likelihood(mvn, X) if self._num_outputs > 1: mean_x = mvn.mean covar_x = mvn.covariance_matrix output_indices = output_indices or range(self._num_outputs) mvns = [ MultivariateNormal( mean_x.select(dim=output_dim_idx, index=t), lazify(covar_x.select(dim=output_dim_idx, index=t)), ) for t in output_indices ] mvn = MultitaskMultivariateNormal.from_independent_mvns( mvns=mvns) posterior = GPyTorchPosterior(mvn=mvn) if hasattr(self, "outcome_transform"): posterior = self.outcome_transform.untransform_posterior(posterior) return posterior