[Question] Batch evaluation of fantasy models

[JackBuck]

🐛 Bug / Question

I am unsure whether this is me misunderstanding how batch dimensions work, or if there is a bug somewhere. I would like to evaluate the mean of a batch of points on a model generated by fantasize. I am using the PosteriorMean class but it is raising errors on a dimension check.

As I understand things, calling .fantasize on a SingleTaskGP creates a batched model, say with batch_dimension = (num_fantasies, 1). If I want to evaluate each fantasy GP on the same batch of inputs, I need to call .posterior with an input of shape b1 x ... x bk x 1 x 1 x m x d where:

• b1 x ... x bk is the batch dimension of the points I wish to evaluate,
• the 1s give space for the batch dimensions of the fantasy model,
• m is the number of points to evaluate jointly,
• d is the dimension of the input space.

I got this understanding from the documentation here: https://botorch.org/docs/batching#batched-evaluation-of-models-and-acquisition-functions

Doing the above experiment, the dimensions do appear to work like this. However, the PosteriorMean class complains about the dimensions through its t_batch_mode_transform decorator and I do not understand why. See the following example.

To reproduce

import torch
from botorch.acquisition import PosteriorMean
from botorch.models import SingleTaskGP
from botorch.sampling import SobolQMCNormalSampler


if __name__ == "__main__":
    train_x = torch.cartesian_prod(  # 20 x 2
        torch.linspace(0, 1, 4),
        torch.linspace(0, 1, 5)
    )
    train_y = torch.randn(20, 1)  # 20 x 1
    model = SingleTaskGP(train_x, train_y)  # batch_shape = []

    fantasy_x = torch.tensor([[[0.5, 0.6]]])  # 1x1x2 (t-batch = 1, q-batch=1)
    sampler = SobolQMCNormalSampler(sample_shape=torch.Size([7]))  # 7 fantasy samples
    fantasy_model = model.fantasize(fantasy_x, sampler) 
    # fantasy_model has batch_shape = 7 x 1 (with the 1 coming from the t-batch shape of the training data)

    # Evaluation batch shape 3x3 => 3x3xmxd where m=1, d=2
    test_x = torch.rand(3, 3, 1, 2)
    # Add in the fantasy batch shape so that they are broadcastable
    fantasy_batch_dim = len(fantasy_model.batch_shape)
    test_x = test_x.reshape(
        test_x.shape[:-2] + fantasy_batch_dim * (1,) + test_x.shape[-2:]
    )
    # test_x now has shape 3x3x1x1xmxd with m=1, d=2
    value_function = PosteriorMean(fantasy_model)
    posterior_mean = value_function(test_x)
    # ^^ This line blows up

AssertionError: Expected the output shape to match either the t-batch shape of X, or the
    `model.batch_shape` in the case of acquisition functions using batch models; but got
    output with shape torch.Size([3, 3, 7, 1]) for X with shape
    torch.Size([3, 3, 1, 1, 1, 2]).

Expected Behavior

The output of shape (3, 3, 7, 1) makes sense to me, since it is what you get when you concatenate (3, 3) (the batch dimension of the inputs) with (7, 1) the batch dimension of the fantasy model. I would expect no error. However, I am unsure whether I have correctly understood all the conventions around batch dimensions here.

System information

Please complete the following information:

• BoTorch Version: 0.8.1
• GPyTorch Version: 1.9.1
• PyTorch Version: 1.13.1+cpu
• Python version: 3.10.9
• Computer OS: Ubuntu 20.04.5 LTS

[saitcakmak]

Hi @JackBuck. Thanks for flagging this and the nice repro. Looks like we check equality of the shapes rather than broadcastability. For now, you can get around the issue by expanding the model batch dimensions.

test_x = test_x.reshape(
    test_x.shape[:-2] + fantasy_batch_dim * (1,) + test_x.shape[-2:]
).expand(test_x.shape[:-2] + fantasy_model.batch_shape + test_x.shape[-2:])
# This will be 3 x 3 x 7 x 1 x 1 x 2

[JackBuck]

Ahh that's what's going on. Thanks - works a treat!

[saitcakmak]

#1715 should fix this for good.