[Bug] _permute_batch has wrong behavior in some instances with LazyEvaluatedKernelTensor

🐛 Bug

In some situations _permute_batch when operating on a LazyEvaluatedKernelTensor that is batched will not properly permute dimensions but do some kind of expand.

To reproduce

I tried hard to find a minimal example, but so far I’ve only been able to reliably reproduce this when using a SAASBO model in Ax. Full repro here: permute_batch_shape_issue_repro.ipynb.txt - this is on current master of gpytorch, botorch, and Ax.

** Stack trace/error message **

The offending part is this one:

ipdb> covar_blocks_lazy.lazy_tensors[0].shape
torch.Size([1, 4, 100, 100])
ipdb> covar_blocks_lazy.lazy_tensors[0]._permute_batch(1, 0).shape
torch.Size([4, 4, 100, 100])

Expected Behavior

Rather than lazy tensor of shape torch.Size([4, 4, 100, 100]) this should return a tensor of size torch.Size([4, 1, 100, 100]), i.e., perform a proper permutation of batch dimensions.

[saitcakmak]

I tried to dig into it a bit with my minimal knowledge of lazy tensors.

https://github.com/cornellius-gp/gpytorch/blob/70eb9f9561bc3ce1f5c00c36e1d865c6b78fca44/gpytorch/lazy/lazy_tensor.py#L190

This line takes in 2 4 x 1 x 100 x s (s takes values from 5, 6, 100) tensors as components, and the kwargs has the kernel and 'last_dim_is_batch': False. The result of the operation has res.shape = 4 x 4 x 100 x 100.

Looking into what happens in res.shape call, it is retrieved from the cache even at the first call. Ignoring that and going into res._size(), it goes into the else part of this check due to kernel.batch_shape=4, which produces the shape 4 x 4 x 100 x 100 (the first expected_shape is 4 x 1 x 100 x 100, lines 250-260 makes it 4 x 4 x 100 x 100):

https://github.com/cornellius-gp/gpytorch/blob/70eb9f9561bc3ce1f5c00c36e1d865c6b78fca44/gpytorch/lazy/lazy_evaluated_kernel_tensor.py#L237-L260

[Balandat]

Hmm yeah interesting idea. I think there might be a few things to be careful about - e.g. we will need to have things like .size() and .shape() return the already permuted shape even if the actual tensor itself hasn’t been permuted yet. But might not be too bad, I’ll give it a try.