Implement tridiagonal solve in numba backend

[ricardoV94]

Also refactors the lapack codegen

---

📚 Documentation preview 📚: https://pytensor--1311.org.readthedocs.build/en/1311/

[codecov]

Codecov Report

Attention: Patch coverage is 50.22624% with 330 lines in your changes missing coverage. Please review.

Project coverage is 81.96%. Comparing base (8a7356c) to head (0757466).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
...tensor/link/numba/dispatch/linalg/solve/general.py	44.54%	61 Missing ⚠️
...nsor/link/numba/dispatch/linalg/solve/symmetric.py	43.33%	51 Missing ⚠️
...ytensor/link/numba/dispatch/linalg/solve/posdef.py	46.98%	44 Missing ⚠️
...or/link/numba/dispatch/linalg/solve/tridiagonal.py	58.25%	43 Missing ⚠️
pytensor/link/numba/dispatch/linalg/utils.py	42.62%	29 Missing and 6 partials ⚠️
...sor/link/numba/dispatch/linalg/solve/triangular.py	40.00%	29 Missing and 1 partial ⚠️
...ensor/link/numba/dispatch/linalg/solve/cholesky.py	43.18%	25 Missing ⚠️
...nk/numba/dispatch/linalg/decomposition/cholesky.py	47.22%	19 Missing ⚠️
pytensor/link/numba/dispatch/linalg/solve/norm.py	61.53%	10 Missing ⚠️
pytensor/link/numba/dispatch/slinalg.py	77.77%	3 Missing and 5 partials ⚠️
... and 1 more

❌ Your patch status has failed because the patch coverage (50.22%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1311      +/-   ##
==========================================
- Coverage   81.98%   81.96%   -0.03%     
==========================================
  Files         188      198      +10     
  Lines       48489    48672     +183     
  Branches     8673     8677       +4     
==========================================
+ Hits        39756    39896     +140     
- Misses       6582     6625      +43     
  Partials     2151     2151              

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/basic.py	78.50% <ø> (ø)
pytensor/link/numba/dispatch/linalg/_LAPACK.py	79.28% <100.00%> (ø)
pytensor/tensor/slinalg.py	93.40% <100.00%> (+0.01%)	⬆️
pytensor/link/numba/dispatch/linalg/solve/utils.py	50.00% <50.00%> (ø)
pytensor/link/numba/dispatch/slinalg.py	69.10% <77.77%> (+25.01%)	⬆️
pytensor/link/numba/dispatch/linalg/solve/norm.py	61.53% <61.53%> (ø)
...nk/numba/dispatch/linalg/decomposition/cholesky.py	47.22% <47.22%> (ø)
...ensor/link/numba/dispatch/linalg/solve/cholesky.py	43.18% <43.18%> (ø)
...sor/link/numba/dispatch/linalg/solve/triangular.py	40.00% <40.00%> (ø)
pytensor/link/numba/dispatch/linalg/utils.py	42.62% <42.62%> (ø)
... and 4 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jessegrabowski]

Looks good, I like the reorganization. We should take it as a template to refactor tensor.linalg

[jessegrabowski]

why cholesky specifically?

[ricardoV94]

It was emitting this warning in a test, strangely the other lapack functions were not? Gotta double check that

[jessegrabowski]

I assume none of them cache. I thought the tests all had a pytest.filterwarnings decorator on them

[ricardoV94]

We should filter the warnings we know of and can't do anything about (even less the users who get them)

[ricardoV94]

So it was the case that the warning was only making it fail with the cholesky, because we had a pytest.mark.filterwarnings("error") in the blockwise test and we tested a blockwise with cholesky. Instead of ignorning I am also making it fail if warnings are emitted on the slinalg test, and thus I am filtering those other Ops here as well.

[jessegrabowski]

nit: this should be split out to a separate function

[ricardoV94]

Sure but numba loves compiling nested functions

[jessegrabowski]

if there's a performance advantage to keeping it inlined, keep it inlined.

I was thinking if you break it out, it fits the template of the other solves with 3 functions: norm, decomp, and solve

[ricardoV94]

I don't love the norm/warnings at all, but that's a separate matter.

I can split, the only concern I could have is compile time but that's not why I inlined it. I'll do a quick benchmark to see if it's warranted

[ricardoV94]

Split it, if anything compile times seem a bit better, probably because at least that inner function can be cached.