Apply a prior on the lengthscale, can't fit likelihood

[defnetanyer]

from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.models.transforms.outcome import Standardize
from botorch.models.transforms.input import Normalize
from botorch.models import SingleTaskGP
from botorch import fit_gpytorch_model

from botorch.test_functions.synthetic import Hartmann

from botorch.optim import optimize_acqf
from botorch.acquisition.analytic import UpperConfidenceBound

from gpytorch.kernels import ScaleKernel
from gpytorch.kernels.matern_kernel import MaternKernel
from gpytorch.priors.torch_priors import GammaPrior

import torch as torch

target_function = Hartmann()

def covar_module(ard_num_dims, batch_shape):
    covar_module = ScaleKernel(
        MaternKernel(
            nu=2.5,
            ard_num_dims=ard_num_dims,
            batch_shape=batch_shape,
            lengthscale_prior= GammaPrior(1.00, 1.00)
        ),
        batch_shape=batch_shape
    )

    return covar_module

def UCB(model, bounds, beta=1.5, q_points=1, num_restarts=20, raw_samples=100):
    UCB = UpperConfidenceBound(model, beta=beta)
    #only one acquisition nothing going in parallel so pretty straightforward optimization problem
    candidates, _ = optimize_acqf(acq_function=UCB, bounds=bounds, q=q_points, num_restarts=num_restarts, raw_samples=raw_samples, options= {"maxiter": 500})
    return candidates    

bounds = torch.tensor([[0.0]*6, [1.0]*6])
train_x = torch.rand((10, 6), dtype = torch.double)

for dimension in range(6):
    lb, ub = bounds[0, dimension].detach(), bounds[1, dimension].detach()
    train_x[:, dimension] = (ub-lb)*train_x[:, dimension] + lb

train_y = target_function(train_x).unsqueeze(-1)

for i in range(10):
    batch_shape = torch.Size([train_y.shape[1]])
    model = SingleTaskGP(train_X=train_x, train_Y=train_y, 
                         input_transform=Normalize(d=train_x.shape[1], bounds=bounds), 
                         outcome_transform=Standardize(m=1), 
                         covar_module = covar_module(ard_num_dims=train_x.shape[1], batch_shape= batch_shape))
    
    #performs the hyperparam fit
    model = model.to(dtype=torch.double)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_model(mll)

    new_x = UCB(model, bounds)
    new_y = target_function(new_x).unsqueeze(-1)

    train_y = torch.cat([train_y, new_y])
    train_x = torch.cat([train_x, new_x])

    print(torch.squeeze(model.covar_module.base_kernel.lengthscale, 0).tolist())

Above is my code where I try to implement a kernel_module with a prior on the lengthscale. SingleTaskGP if not given a covar_module uses a Matern5/2 kernel with prior on lengthscale as I recreated. I get the error below when I call
fit_gpytorch_model(mll) .

The error is this: RuntimeError: permute(sparse_coo): number of dimensions in the tensor input does not match the length of the desired ordering of dimensions i.e. input.dim() = 1 is not equal to len(dims) = 3

Would anyone be able to help out with this? I'm guessing what the sparse_coo is referring to the coordinate in covariance matrix since the columns are computed as lazy tensors and the covar matrix is a sparse matrix. That doesn't help me undestand the problem, why is it applying a permutation on it? What are the ordering dimensions? And why does my input tensor only have 1 dimension when it should have 6?

Anything helps, thanks everyone!

[Balandat]

Hi, I think there is some confusion about the batch_shape on your end. In this case you don't actually have a batched model, and so you don't have to set a batch_shape (functionally this means it's set to torch.Size([]) on the backend). If you change the line

covar_module = covar_module(ard_num_dims=train_x.shape[1], batch_shape= batch_shape))

to

covar_module = covar_module(ard_num_dims=train_x.shape[1], batch_shape=torch.Size()))

then this runs fine.

Now there is one thing to note - in the multi-output case (when you have more than one trailing dimension in train_y, then the model does some reshuffling of dimensions internally to fit these models as batched models for efficiency/speed reasons. In that case you'll need to use the _aug_batch_shape ("augmented batch shape" property of the model) as done here. But this is not necessary in your current setup.

[defnetanyer]

Thank you so much! Runs fine now.