Add LAPACK overloads for all variants of pt.linalg.solve

[jessegrabowski]

Description

Goal of this PR is to give numba mode full coverage of scipy.linalg.solve options. We currently only support assume_a = "gen". If users select a different solver, they get incorrect results (see #422 ). This PR should fix that bug, plus add:

• support for overwrite_a in numba mode
• support for overwrite_b in numba mode
• support for transposed argument (all modes)
• lu_factor and lu_solve Ops (all modes)
• support for assume_a = "sym" and assume_a = "pos" in numba mode
• support for cho_solve in numba mode

We get the lu_factor and lu_solve Ops kind of "for free" because I'm adding overloads for dgetrs and dgetrf. We just have to write the Ops and do the JAX dispatch. JVP for lu_factor is here. Help wanted. I might decide that these Ops are out of scope for this PR and open another one.

Related Issue

• Closes BUG: pt.linalg.solve returns incorrect results when mode = "NUMBA" #422

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[codecov]

Codecov Report

Attention: Patch coverage is 52.49042% with 248 lines in your changes missing coverage. Please review.

Project coverage is 81.87%. Comparing base (4e85676) to head (fe97e5d).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/slinalg.py	42.62%	205 Missing and 5 partials ⚠️
pytensor/link/numba/dispatch/_LAPACK.py	75.32%	32 Missing and 6 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1146      +/-   ##
==========================================
- Coverage   82.14%   81.87%   -0.28%     
==========================================
  Files         186      187       +1     
  Lines       48210    48617     +407     
  Branches     8678     8705      +27     
==========================================
+ Hits        39603    39804     +201     
- Misses       6440     6640     +200     
- Partials     2167     2173       +6     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/basic.py	80.37% <ø> (-1.86%)	⬇️
pytensor/tensor/slinalg.py	93.49% <100.00%> (+0.01%)	⬆️
pytensor/link/numba/dispatch/_LAPACK.py	75.32% <75.32%> (ø)
pytensor/link/numba/dispatch/slinalg.py	44.42% <42.62%> (-8.38%)	⬇️

[jessegrabowski]

This is pretty close. I just need some help with the destroy_map stuff on Solve. I guess this wasn't being before? The code is a bit hard to follow with all the subclassing.

[ricardoV94]

What do you mean we were not doing inplace? It was already implemented for the default backend for Solve and Cholesky: ed6ca16#diff-f5e942a523e3c0402aa63824184c381f5a963867125a21f119e2dab97a110d95

Just not in Numba yet. You shouldn't change anything in how the Ops are created, it's handled by the inplace rewrites. You may need to explicitly trigger them in the numba tests (the default numba mode in compare_py_and_numba does not include them), but I am not sure that utility actually works with inplace stuff

[ricardoV94]

Doesn't this require changes in the grads?

[jessegrabowski]

Yes, good catch. There's some logic in the SolveBase L_Op that I thought was handling this, but it's actually handling the lower flag for SolveTriangular and ChoSolve I think. It should be similar/the same, though?

At any rate I'll add some gradient tests to make sure