Gpytorch: [Bug] Sampling from priors doesn't match shape of hyperparameters

[mshvartsman]

🐛 Bug

I found some unexpected interactions between ard_num_dims and the shapes of priors for kernels -- a few settings where if I sample from a hyperparameter prior I don't get a tensor the same shape as the hyperparameter. I'm not sure if all of these are intended or not, but looks like a bug to me.

To reproduce

import torch
from gpytorch.priors import GammaPrior, NormalPrior
from gpytorch.kernels import RBFKernel

# make a kernel
scales = torch.Tensor([1,1])
kernel = RBFKernel(
    ard_num_dims=2,
    lengthscale_prior=GammaPrior(3.0, 6.0 / scales),
)
new_lengthscale = kernel.lengthscale_prior.sample(kernel.lengthscale.shape)
print(kernel.lengthscale.shape) # size 1,2
print(new_lengthscale.shape) # size 1,2,2, if I try to assign it back I get an error

# same with another prior
kernel2 = RBFKernel(
    ard_num_dims=2,
    lengthscale_prior=NormalPrior(loc=10, scale=scales)
)

new_lengthscale = kernel2.lengthscale_prior.sample(kernel2.lengthscale.shape)
print(kernel2.lengthscale.shape) # size 1,2
print(new_lengthscale.shape) # size 1, 2, 2

# ard_num_dims is only 1 but we have a higher-dim prior. Is this behavior defined?
kernel3 = RBFKernel(
    ard_num_dims=1,
    lengthscale_prior=NormalPrior(loc=10, scale=scales)
)

new_lengthscale = kernel3.lengthscale_prior.sample(kernel3.lengthscale.shape)
print(kernel3.lengthscale.shape) # size 1, 1 -- but shouldn't we expect 1,2? 
print(new_lengthscale.shape) # size 1, 1, 2

# ok, ard_num_dims is 2 but my prior is 1d, now it works correctly
kernel4 = RBFKernel(
    ard_num_dims=2,
    lengthscale_prior=NormalPrior(loc=10, scale=1)
)

new_lengthscale = kernel4.lengthscale_prior.sample(kernel4.lengthscale.shape)
print(kernel4.lengthscale.shape) # size 1,2
print(new_lengthscale.shape) # size 1, 2

Expected Behavior

It would be nice if we got a warning/error earlier for undefined/unsupported behavior, and otherwise shapes matched correctly.

System information

Please complete the following information:

• GPyTorch Version: 1.2.0
• PyTorch Version: 1.6.0.
• Computer OS: verified on OSX and CentOS.

[Balandat]

Yeah @dme65 ran into some issues that are likely related, it seems like there are some issues with batch sizes and the priors.

[mshvartsman]

More thoughts about this: priors are just thin wrappers around pytorch distributions, right? The sample method on those

Generates a sample_shape shaped sample or sample_shape shaped batch of samples if the distribution parameters are batched.
(https://pytorch.org/docs/stable/distributions.html#torch.distributions.distribution.Distribution.sample)

So if my prior is not batched, then in order to get the right size of kernel hyperparams I need to prior.sample(hyperparam.shape) but if my prior is batched with batch size hyperparam.shape (or multivariate) I need to prior.sample(torch.Size([1])), which is awkward.

So, what should we do about it? Some options:

• Option 1: ensure consistency in the kernel constructor. Something like the following:

    If prior is univariate with batch_size == 1: 
        promote prior to batched with batch_size=hyperparam.shape
    elif prior is multivariate or univariate with nonzero batch_size: 
        check that prior size (multivariate) or batch size (univariate) matches hyperparam.shape

(and maybe, warn on calls to prior.sample(...) with argument equal to the prior batch size). There are probably reasons to call prior.sample() with various sizes but so this might be an overly stringent warning, but it'd reduce this confusion in the future.

• Option 2: enforce consistency but don't change the priors. So basically, if prior.sample(1) doesn't generate something of lengthscale.shape, throw an error. .
• Option 3: rely on client code to sample more intelligently from the priors (in particular, if no changes are made on gpytorch side I think code like this https://github.com/pytorch/botorch/blob/754f5ea32858821d83d8199d375865bceb395085/botorch/optim/utils.py#L31 probably needs to be changed)

Thoughts? Happy to PR this if there's consensus on how to proceed.

[Balandat]

So Option 1 makes sense to me. We could just use the existing expand method to achieve this promotion.

I think one potential challenge with this is that since this will not be in-place if we pass a prior as a prior to a module constructor, then the registered prior will refer to a different object. Specifically, say prior is a non-batched prior, then

prior.sample(torch.Size([1]))  # size 1
k = Kernel(batch_shape=torch.Size([2]), lengthscale_prior=prior)
k._priors["lengthscale_prior"].sample(torch.Size([1]))  # size 2 x 1

This is probably ok though...

A similar shape promotion would have to be done for ARD I guess.

[mshvartsman]

Could we override prior.sample to do the checks rather than do it on the __init__ and call expand at the last possible moment? Then the object is the same, though the size mismatch in your example will still come up, and we're eating all the extra memcopy costs, so I guess that's not better unless we also do some caching, at which point they're different objects again.

[Balandat]

Yeah we could do that. I guess the one concern would be perf if we repeatedly sample from the prior and have to do the expansion each time. This shouldn't be too much of an issue though unless we want to sample a ton, e.g. when doing fully Bayesian inference (cc @jpchen)

On the other hand, I don't think too many folks will do something like k._priors["lengthscale_prior"].sample, so the point I brought up might not really be a valid concern.