[Bug]: MVNXPB doesn't support multiple batch dimensions

[JackBuck]

What happened?

The documentation for MVNXPB indicates that covariance_matrix and bounds accept a batch_shape. However, when called with a batch_shape of more than one dimension it blows up. I would expect it to give the same as calling the algorithm in a loop for each element of the batch. Alternatively, the docs could specify only that one batch dimension is allowed, and the code should raise an informative error if more than one is passed.
https://botorch.readthedocs.io/en/stable/utils.html#botorch.utils.probability.mvnxpb.MVNXPB

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

import torch
from botorch.utils.probability import MVNXPB

# Create inputs with a batch_shape of (2,2)
covmat = torch.diag(torch.tensor([1.0, 2.0, 3.0])**2).expand(2, 2, -1, -1)
bounds = torch.tensor(
    [
        [
            [[0, torch.inf], [0, torch.inf], [0, torch.inf]],
            [[-torch.inf, 0], [-torch.inf, 0], [-torch.inf, 0]],
        ],
        [
            [[1, torch.inf], [1, torch.inf], [1, torch.inf]],
            [[-torch.inf, -1], [-torch.inf, -1], [-torch.inf, -1]],
        ],
    ]
)

# Create and call the solver
solver = MVNXPB(covmat, bounds)
log_probs = solver.solve()

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

IndexError                                Traceback (most recent call last)
Cell In[1], line 21
     19 # Create and call the solver
     20 solver = MVNXPB(covmat, bounds)
---> 21 log_probs = solver.solve()

File .../lib/python3.13/site-packages/botorch/utils/probability/mvnxpb.py:236, in MVNXPB.solve(self, num_steps, eps)
    234     self.step += 1
    235     if should_update_chol:
--> 236         piv_chol.update_(eps=eps)
    238 # Factor in univariate probability if final term fell outside of a block.
    239 if self.step % 2:

File .../lib/python3.13/site-packages/botorch/utils/probability/linalg.py:128, in PivotedCholesky.update_(self, eps)
    126 rank1 = (rank1 * rank1.transpose(-1, -2)).tril()
    127 L[..., i + 1 :, i + 1 :] = L[..., i + 1 :, i + 1 :].clone() - rank1
--> 128 L[Lii <= i * eps, i:, i] = 0  # numerical stability clause
    129 self.step += 1

IndexError: The shape of the mask [2, 2] at index 1 does not match the shape of the indexed tensor [2, 1, 3] at index 1

BoTorch Version

0.13.0

Python Version

3.13.2

Operating System

Ubuntu 20.04.6 LTS (Focal Fossa)

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[ltiao]

Thank you for bringing this issue to our attention. We appreciate the detailed report and the reproducible example you've provided.

It appears that the MVNXPB method does indeed encounter issues when len(batch_shape) > 1. In particular, the line where the exception is raised does not adequately handle this case. We will work on fixing this as soon as possible and will also add additional test cases to ensure this scenario is properly supported in the future.

[JackBuck]

Thanks Louis!

[ltiao]

If you need to unblock yourself urgently, it appears that updating the relevant line in PivotedCholesky.update_ as follows fixes the issue:

L[..., i:, i][Lii <= i * eps] = 0  # numerical stability clause

I will put up a pull request with this fix shortly. In the meantime, you can also move forward with this fix on a local/editable installation of BoTorch.

Thank you again for pointing out this issue!

[JackBuck]

Thanks @ltiao !