Fix bug in tag_solve_triangular rewrite

[jessegrabowski]

Motivation for these changes

Currently, the tag_solve_trangular rewrite looks for graphs like this:

import pytensor
import pytensor.tensor as pt
import numpy as np
from numpy.testing import assert_allclose

A = pt.dmatrix('A')
b = pt.dmatrix('b')
L = pt.linalg.cholesky(A)
L_inv = pt.linalg.solve(L, b, assume_a='gen')
f = pytensor.function([A, b], [L_inv])

pytensor.dprint(L_inv)

Solve{assume_a='gen', lower=True, check_finite=True} [id A]
 ├─ Cholesky{lower=True, destructive=False, on_error='raise'} [id B]
 │  └─ A [id C]
 └─ b [id D]

And replaces then with this:

pytensor.dprint(f)

Solve{assume_a='sym', lower=True, check_finite=True} [id A] 1
 ├─ Cholesky{lower=True, destructive=False, on_error='raise'} [id B] 0
 │  └─ A [id C]
 └─ b [id D]

The Solve(assume_a='sym', lower=True) solver was incorrectly assumed to be a triangular solver. In fact, it is a symmetric solver that decomposes the symmetric A matrix into a triangular matrix. But given a triangular matrix, it produces incorrect results. This can be verified directly by testing L @ L_inv = eye:

n = 3
Z = np.random.normal(size=(n, n))
P = Z @ Z.T
P_chol = np.linalg.cholesky(P)
eye = np.eye(n)

assert_allclose(f(P, eye) @ P_chol, eye) # fails

Actually, even before the rewrite this code fails, so there's something deeper going on:

assert_allclose(f(L_inv.eval({A:P, b:eye}) @ P_chol, eye) # also fails

But one step at a time. This PR replaces Solve(assume_a='sym', lower=True) with SolveTriangular in the tag_solve_triangular rewrite, and adds a test for correct computation to the rewrite. That computational test can be put somewhere else once we track down what's going on with Solve more generally, but it's a good stop-gap for now.

Checklist

• Explain motivation and implementation 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

New features

• ...

Bugfixes

• Graphs with a Cholesky Op followed by a Solve op will no longer be incorrectly computed.

Documentation

• ...

Maintenance

• ...

Closes #382

[ricardoV94]

WDYT? The test name should also be changed then

Suggested change

	def tag_solve_triangular(fgraph, node):
	def solve_cholesky_to_solve_triangular(fgraph, node):

[jessegrabowski]

Yeah agreed, "tag" doesn't mean anything. But we shouldn't call it "solve_cholesky", because there's a separate Op called SolveCholesky. Maybe generic_solve_to_solve_triangular?

[ricardoV94]

Sounds good

[ricardoV94]

You got one of those pesky float32 test failures. LGTM otherwise. Just left a suggestion of giving a more descriptive name to the rewrite

[codecov-commenter]

Codecov Report

Merging #383 (455707e) into main (7a82a3f) will decrease coverage by 0.01%.
The diff coverage is 69.23%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #383      +/-   ##
==========================================
- Coverage   80.46%   80.46%   -0.01%     
==========================================
  Files         156      156              
  Lines       45515    45514       -1     
  Branches    11149    11148       -1     
==========================================
- Hits        36625    36624       -1     
  Misses       6688     6688              
  Partials     2202     2202              

Impacted Files	Coverage Δ
pytensor/tensor/rewriting/linalg.py	71.96% <69.23%> (-0.26%)	⬇️