[Bug] condition_on_observation with SaasFullyBayesianGP

[hvarfner]

🐛 Bug

With the risk of looking a bit silly, I don't think conditioning on additional observations works when using fully Bayesian GPs out of the box.

To reproduce

The SAASBO tutorial works fine to reproduce. ** Code snippet to reproduce **

import os

import matplotlib.pyplot as plt
import numpy as np
import torch
from botorch import fit_fully_bayesian_model_nuts
from botorch.acquisition import qExpectedImprovement
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
from botorch.models.transforms import Standardize
from botorch.optim import optimize_acqf
from botorch.test_functions import Branin
from torch.quasirandom import SobolEngine

SMOKE_TEST = os.environ.get("SMOKE_TEST")
tkwargs = {"device": torch.device("cuda:1" if torch.cuda.is_available() else "cpu"), "dtype": torch.double}

WARMUP_STEPS = 512 if not SMOKE_TEST else 32
NUM_SAMPLES = 256 if not SMOKE_TEST else 16
THINNING = 16

train_X = torch.rand(10, 4, **tkwargs)
test_X = torch.rand(5, 4, **tkwargs)
train_Y = torch.sin(train_X[:, :1])
test_Y = torch.sin(test_X[:, :1])

gp = SaasFullyBayesianSingleTaskGP(train_X=train_X, train_Y=train_Y)
fit_fully_bayesian_model_nuts(
    gp, warmup_steps=WARMUP_STEPS, num_samples=NUM_SAMPLES, thinning=THINNING, disable_progbar=True
)
with torch.no_grad():
    posterior = gp.posterior(test_X)

# Shape check
print(gp.covar_module.base_kernel.lengthscale.shape) # 16 x 1 x 4

cond_X = torch.rand(5, 4, **tkwargs) # 5 x 4
cond_Y = torch.sin(cond_X[:, :1]) # 5 x 1

# RuntimeError
gp.condition_on_observations(cond_X, cond_Y)

# and accounting for the batch shape does not work, either
cond_X = torch.rand(5, 4, **tkwargs).repeat(16, 1, 1) # 16 x 5 x 4
cond_Y = torch.sin(cond_X[:, :1]).repeat(16, 1, 1) # 16 x 5 x 1
gp.condition_on_observations(cond_X, cond_Y)
# Please make sure it does not require any external dependencies

** Stack trace/error message ** For both cases, this occurs:

gpytorch/models/exact_prediction_strategies.py:131, in DefaultPredictionStrategy.get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output, **kwargs)
    127 full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
    129 batch_shape = full_inputs[0].shape[:-2]
--> 131 full_mean = full_mean.view(*batch_shape, -1)
    132 num_train = self.num_train
    134 # Evaluate fant x train and fant x fant covariance matrices, leave train x train unevaluated.

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

Expected Behavior

Conditioning on the new observations for every constituent GP in the FBGB.

My two cents on the issue

Notably, creating a BatchedMultiOutputModel with num_mcmc_samples batches of train inputs and targets handles this without issue. The issue seems to appear here: https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/models/exact_gp.py#L220, where the output is of shape num_mcmc_samples x num_mcmc_samples x num_inputs (should probably be num_mcmc_samples x num_inputs.

System information

Please complete the following information:

• Ubuntu 18.04