[Bug]: RuntimeError: Index out of bounds

[AdrianSosic]

What happened?

I have snippet of code that works without problems on macOS but throws a RuntimeError on Linux/Windows.
So far, I haven't fully figured out what is going on, but what I observed:

• The MultiTaskGP.task_mapper is assigned [nan, 0]
• The above is perhaps still expected, but in GPyTorch, these nan values get translated into 9223372036854775808, i.e. some integer overflow is happening (the number is the -2ˆ63)

Please provide a minimal, reproducible example of the unexpected behavior.

import torch
from botorch.acquisition import qLogExpectedImprovement
from botorch.acquisition.objective import IdentityMCObjective
from botorch.models import MultiTaskGP
from botorch.optim import optimize_acqf_mixed

train_x = torch.tensor([[1.0, 10.0], [1.0, 20.0]])
train_y = torch.tensor([[10.0], [20.0]])
gp = MultiTaskGP(
    train_x,
    train_y,
    task_feature=0,
    output_tasks=[0, 1],
    all_tasks=[0, 1],
    validate_task_values=False,
)


acqf = qLogExpectedImprovement(gp, 0.0, objective=IdentityMCObjective())
bounds = torch.tensor([[0.0, 0.0], [1.0, 5.0]])
q = 1
num_restarts = 10
raw_samples = 64
fixed_features_list = [{0: 0}, {0: 1}]

candidate, acqf_value = optimize_acqf_mixed(
    acqf,
    bounds=bounds,
    q=q,
    num_restarts=num_restarts,
    raw_samples=raw_samples,
    fixed_features_list=fixed_features_list,
)

print(candidate)
print(acqf_value)

Please paste any relevant traceback/logs produced by the example provided.

RuntimeError: index -9223372036854775808 is out of bounds for dimension 2 with size 1

BoTorch Version

0.16.1

Python Version

3.10.19

Operating System

Linux/Windows

(Optional) Describe any potential fixes you've considered to the issue outlined above.

No response

Pull Request

None

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[Balandat]

Hmm interesting. Assigning nan is indeed expected here. Can't repro this on MacOS either; but let me try to repro this on Linux to understand how exactly this propagates. cc @sdaulton who has worked on this task mapping logic.

[Balandat]

@AdrianSosic I was going to ask why you want to predict on tasks that you don't have training data for, but given your discssion on #3085 it seems that you indeed intend to do this.

Looking into this a bit more, the operation in which we end up with the overflow is the long() conversion here:

botorch/botorch/models/multitask.py

Line 114 in abb09b6

mean_x = mean_x.gather(-1, task_idcs.long()).squeeze(-1)

This change was introduced by @hvarfner in #3049, which made the MultiTaskGP model default to a MultiTaskMean instead of a regular Mean. My guess was initially that this has worked in the past when using a regular mean function, but when passing that in I run into other issues with overflow indices in torch.dsmm deep inside linear_operator.

Before I get much deeper into this rabbit hole - @sdaulton. @hvarfner what do you think should be the proper behavior in this case? I'd like to avoid having to special case all kinds of things regarding nan indices in botorch and gpytorch, so ideally what we want to do here would not require that.