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
