Skip to content

Fix CholLinearOperator with upper triangular factors and its slicing - #145

Open
M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:fix/chol-upper-and-getitem
Open

M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:fix/chol-upper-and-getitem

Conversation

@M-Colley

Copy link
Copy Markdown

Fixes #144.

Problem

CholLinearOperator inherits most of its methods from RootLinearOperator, which assumes that the matrix is root @ root.mT (see #144 for reproductions):

  1. For an upper triangular factor (upper=True, i.e. R^T R), matmul, diagonal, indexing, expand, constant multiplication, root_decomposition and sampling used R R^T instead. Also, upper was not part of the representation, so detach(), clone(), to() and the autograd functions rebuilt the operator with upper=False.
  2. inverse() returned a CholLinearOperator with upper=True whose factor L^{-1} is lower triangular, so its matmul, diagonal and solve were wrong.
  3. Slicing with equal row and column indices, e.g. op[0:3, 0:3], constructed a CholLinearOperator from the selected rows of the factor, which are not triangular, and raised NotImplementedError: torch.tril(InterpolatedLinearOperator, ) is not implemented.

Change

  1. Store an upper triangular factor R as the equivalent lower triangular factor L = R^T (a lazy transpose of the TriangularLinearOperator). Then all inherited methods and the reconstruction from the representation are correct. upper is still accepted, but the upper attribute is now always False and root is R^T. The if self.upper: branches in to_dense, inv_quad etc. are kept, although they are no longer taken; I can remove them if you prefer.
  2. inverse() returns RootLinearOperator(L^{-T}), since (L L^T)^{-1} = L^{-T} L^{-1} is not a Cholesky factorization. This is the same representation that root_inv_decomposition() already returns.
  3. CholLinearOperator._getitem returns a RootLinearOperator of the selected rows for indices of the matrix dimensions. Batch indexing still returns a CholLinearOperator.

Tests

  • TestCholLinearOperatorUpper, ...UpperBatch and ...UpperMultiBatch run the whole LinearOperatorTestCase suite for upper triangular factors. Their reference matrix uses the operator's upper attribute, so the tests also catch upper being lost on reconstruction. On main, 26 different tests fail for them, including test_getitem, test_diagonal, test_sample, the test_solve_* tests and test_constant_mul.
  • test_getitem_square_slices slices the rows and columns with the same index; it fails on main for all 6 classes.
  • test_inverse also checks matmul and solve of the inverse; it fails on main for all 6 classes. Its gradient comparison now skips non-leaf representation tensors (the lazily transposed factor of the upper variants), like the generic tests do.

Full suite: 5238 tests pass, except test_psd_safe_cholesky_psd, which also fails on main locally (torch 2.13 on Windows, unrelated). flake8 passes. Downstream, BoTorch's tests that use Cholesky covariances give the same results with and without this change, and BoTorch's separate_mtmvn works for CholLinearOperator covariances without a workaround. Tested with Python 3.13, PyTorch 2.13 (CPU).

🤖 Generated with Claude Code

CholLinearOperator inherits most of its methods from RootLinearOperator,
which assumes that the matrix is root @ root.mT. For an upper triangular
factor (upper=True, i.e. R^T R), matmul, diagonal, indexing, expand,
constant multiplication, root_decomposition and sampling therefore used
R R^T instead. Also, upper was not part of the representation, so detach,
clone, to() and the autograd functions rebuilt the operator with
upper=False. Store an upper triangular factor R as the equivalent lower
triangular factor R^T instead, so that all of these methods are correct.

inverse() returned a CholLinearOperator with upper=True whose factor
L^{-1} is lower triangular, so its matmul, diagonal and solve were wrong.
(L L^T)^{-1} = L^{-T} L^{-1} is not a Cholesky factorization, so return a
RootLinearOperator of L^{-T} instead.

Slicing the rows and columns of a CholLinearOperator with the same index,
e.g. op[0:3, 0:3], constructed a CholLinearOperator from the selected
rows of the Cholesky factor, which are not triangular, and raised
"NotImplementedError: torch.tril(InterpolatedLinearOperator, ) is not
implemented". Return a RootLinearOperator of these rows instead. Batch
indexing still returns a CholLinearOperator.

Run the CholLinearOperator tests for upper triangular factors too, and
test slicing with equal row and column indices and the matmul and solve
of the inverse.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] CholLinearOperator: wrong values for upper triangular factors and inverse(), slicing raises

1 participant