Fix mis-broadcast in exact_predictive_mean for batched multi-output GPs

[saitcakmak]

Summary

DefaultPredictionStrategy.exact_predictive_mean unconditionally squeezed dim 1 of any 4-dimensional mean_cache. This is only correct for multitask models (MultitaskMultivariateNormal), whose mean_cache carries a spurious singleton at index 1. For batched multi-output (batch-independent) models — a plain, non-multitask MultivariateNormal — a 4-dim mean_cache with a singleton batch dimension gets its batch dim dropped, and the remaining dimensions mis-broadcast against test_train_covar.

Fixes the crash reported in meta-pytorch/botorch#3336 ("Fantasizing on Multioutput Batch Independent GP adds another output dimension").

Root cause

The len == 4 shape heuristic was introduced in commit b9dc064 ("Passing unit tests now", part of the original qKG work), replacing the earlier isinstance(train_prior_dist, MultitaskMultivariateNormal) gating. The shape-only check over-fires: it also catches plain batched multi-output models, for which the squeeze is wrong.

The bug only triggers when dim 1 (the candidate batch b) has size 1. For b > 1, squeeze(1) is a no-op, which is why single-output GPs and q-batched cases were unaffected.

Shape trace

For a fantasized batched multi-output GP with a single candidate (b == 1):

mean_cache        [8, 1, 2, 12]  --squeeze(1)-->  [8, 2, 12]
test_train_covar  [64, 8, 1, 2, 21, 12]

The squeezed cache mis-broadcasts, multiplying the two num_fantasies=8 dims into [64, 8, 8, 2, 21] (numel 172032) instead of aligning into [64, 8, 1, 2, 21] (numel 21504). predictive_mean.view(...) in ExactGP.__call__ then fails:

RuntimeError: shape '[64, 8, 1, 2, 21]' is invalid for input of size 172032

Fix

Guard the squeeze with the distribution-type check the code originally had, restoring the pre-b9dc064 gating. Multitask behavior is byte-for-byte identical.

Minimal repro (pure BoTorch)

import torch
from botorch.models import SingleTaskGP
b0, b1, n, d, m = 8, 1, 11, 2, 2
tX = torch.rand(b0, b1, n, d, dtype=torch.double)
tY = torch.rand(b0, b1, n, m, dtype=torch.double)
model = SingleTaskGP(tX, tY)
model.posterior(torch.rand(64, b0, b1, 21, d, dtype=torch.double)).mean
# unpatched: RuntimeError: shape '[64, 8, 1, 2, 21]' is invalid for input of size 172032
# patched:   returns mean of shape [64, 8, 1, 21, 2]

Testing

Added test_batched_multioutput_prediction_singleton_batch in test/models/test_exact_gp.py. It reproduces the crash deterministically (no fantasize randomness) using a batched multi-output GP whose training data has a 3-dim aug batch [b0, b1, m] = [4, 1, 2] (producing the same 4-dim mean_cache), and checks the multi-output means against independent per-output single-output GPs to < 1e-6. Mutation-checked: the test fails on unpatched code (RuntimeError: shape '[64, 4, 1, 2, 5]' is invalid for input of size 10240) and passes after the fix.

Verification:

• New regression test: fails unpatched, passes patched. ✅
• GPyTorch exact-GP + multitask + fantasy suites (test_exact_gp, test_simple_gp_regression, test_batch_multitask_gp_regression, test_independent_multitask_gp_regression, test_kronecker_multitask_gp_regression, test_hadamard_multitask_gp_regression, test_derivative_gp_fantasy): 88 passed (85 baseline + 3 inherited runs of the new test). ✅
• Downstream BoTorch smoke tests (test_gpytorch, test_gp_regression, test_model_list_gp_regression, test_knowledge_gradient): 69 passed. ✅
• flake8 + ufmt clean. ✅