MVN sampling of linear mean parameters

[stevenasheets]

I'm trying to set a MVN prior on the linear mean parameters, basically transferring the posterior distribution of a bayesian linear regression to the linear mean function of a GP.

Right now I'm following the code examples given in the Fully Bayesian GP example.

My code is:

# We will use the simplest form of GP model, exact inference
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = LinearMean(train_x.shape[-1])
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

num_samples = 100
warmup_steps = 100
N,D = x.shape

mean_prior_vec = torch.zeros(D)
mean_prior_cov = torch.eye(D)

likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.Positive())
model = ExactGPModel(x, y, likelihood)
model.mean_module.register_prior("mean_prior_weights", MultivariateNormalPrior(mean_prior_vec, mean_prior_cov), "weights")
model.mean_module.register_prior("mean_prior_bias", NormalPrior(0, 1), "bias")
model.covar_module.base_kernel.register_prior("lengthscale_prior", UniformPrior(0.01, 0.5), "lengthscale")
model.covar_module.register_prior("outputscale_prior", UniformPrior(1, 2), "outputscale")

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

def pyro_model(x, y):
    with gpytorch.settings.fast_computations(False, False, False):
        sampled_model = model.pyro_sample_from_prior()
        output = sampled_model.likelihood(sampled_model(x))
        pyro.sample("obs", output, obs=y)
    return y

nuts_kernel = NUTS(pyro_model)
mcmc_run = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps)
mcmc_run.run(x, y)

This results in a runtime error:

File ~/anaconda3/lib/python3.9/site-packages/gpytorch/priors/torch_priors.py:137, in MultivariateNormalPrior.expand(self, batch_shape)
    135 batch_shape = torch.Size(batch_shape)
    136 cov_shape = batch_shape + self.event_shape
--> 137 new_loc = self.loc.expand(batch_shape)
    138 new_scale_tril = self.scale_tril.expand(cov_shape)
    140 return MultivariateNormalPrior(loc=new_loc, scale_tril=new_scale_tril)

RuntimeError: The expanded size of the tensor (1) must match the existing size (3) at non-singleton dimension 1.  Target sizes: [3, 1].  Tensor sizes: [3]
Trace Shapes:
 Param Sites:
Sample Sites:

I'm stumped trying to figure out how to bypass this error. If I switch to a NormalPrior on the linear mean weights it will work but I'd prefer to do the full covariance transfer if possible. Any ideas? Thanks.