Allow transposed argument in linalg.solve

[jessegrabowski]

Description

Adds a transposed argument to pt.linalg.solve. The signature now matches scipy.linalg.solve 1-to-1. It also fixes gradients for pt.linalg.solve_triangular when trans is either "T" or "C" (though we don't really support complex inputs anyway)

It also adds a bunch of tests for the two implicated solve Ops.

Related Issue

• Closes Add transposed argument to pt.linalg.solve #1229
• Related to #

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--1231.org.readthedocs.build/en/1231/

[ricardoV94]

question does this matter for PyTensor? what would we lose if the transpose flag was just used symbolically but not part of the Op? Say if the user calls solve(transposed=True) we transpose the input and call the old op without the transpose flag?

transpose should be pretty cheap and if you have 2 in the grad they get rewritten away anyway.

Is this not a concern only for eager libraries like scipy? I wouldn't be surprised if Jax also does it just symbolically

[jessegrabowski]

I have no objection to going that way instead. It might be smarter -- It seems like the transposed flag might be a bit slower in pure numpy:

from scipy import linalg
import numpy as np
rng = np.random.default_rng()

A = rng.normal(size=(10_000, 10_000))
b = rng.normal(size=(10_000,))

def f1(A, b):
    return linalg.solve(A.T, b)

def f2(A, b):
    return linalg.solve(A, b, transposed=True)

%timeit f1(A, b)
# 1.84 s ± 107 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f2(A, b)
# 2.37 s ± 82.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Plus we have to transpose the output of the gradients when transposed=True anyway, so there's still a transpose popping up somewhere. I'll make the change to just have it quietly transpose the inputs for you.

The bottom line is that I'd like for the argument to be there so we match the scipy API.

[ricardoV94]

The bottom line is that I'd like for the argument to be there so we match the scipy API.

Yeah we should always try to match the external API

[ricardoV94]

Are all the test changes needed? At least the transpose in JAX seems unnecessary since it's not part of the op

[ricardoV94]

You're testing 36 combinations with parametrize just for solve in jax, transpiling both python and JAX and jitting/evaluating the latter. Doesn't seem reasonable

[jessegrabowski]

But it's fast

[ricardoV94]

But it's fast

That's true. Maybe you should add dtype parametrization to validate my concerns :)

[jessegrabowski]

brb adding those tests

In other news, why is jax pad test failing? Is it flakey?

[ricardoV94]

Hmm I saw it failing in my other PR and I thought it was my changes. In that case either something broke in a recent jax release or a commit we merged into main is causing it to fail

[maresb]

Jax 0.4.36 just got released on conda-forge about 9 hours ago, so maybe something changed with that. You could test if pinning jax <0.4.36 works around the issue. conda-forge/jax-feedstock#169

[ricardoV94]

I tested locally jax had a bug in 0.4.36 that seems to have disappeared by 0.5.1. I didn't bother checking exactly when did it stop failing. Added a commit in another PR that will skip it: a0d9ecf

[ricardoV94]

Nope it's still failing: jax-ml/jax#26888

[ricardoV94]

you can rebase @jessegrabowski

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.99%. Comparing base (757a10c) to head (0f26954).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1231   +/-   ##
=======================================
  Coverage   81.99%   81.99%           
=======================================
  Files         188      188           
  Lines       48600    48608    +8     
  Branches     8685     8688    +3     
=======================================
+ Hits        39849    39857    +8     
  Misses       6586     6586           
  Partials     2165     2165           

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/slinalg.py	95.00% <ø> (-0.13%)	⬇️
pytensor/link/numba/dispatch/slinalg.py	44.42% <ø> (-0.11%)	⬇️
pytensor/tensor/slinalg.py	93.67% <100.00%> (+0.13%)	⬆️