Prevent the recalculation of permutation in lu_solve()

[ahoenselaar]

Use case: Repeated solution of a linear system based on LU decomposition.

Currently, the only exposed methods to make this possible are jax.scipy.linalg.lu_factor() in conjuction with jax.scipy.linalg.lu_solve(). However, to conform with SciPy's API, jax.scipy.linalg.lu_factor() drops the third output permutation provided by jax.lax.linalg.lu(). Hence, the permuation is recalculated in every invocation of jax.scipy.linalg.lu_solve() (code).

For reasons that should be investigated separately, lu_pivots_to_permutation() exhibits poor performance on GPU.

There are two possible solutions:

• Add an optional argument to jax.scipy.linalg.lu_solve() that allows the user to pass in permutation. That would cause divergence from SciPy's API.
• Expose jax.lax.linalg.lu_solve() directly to the user.

[hawkinsp]

I'm curious if you think this is still worth solving, given that #5880 significantly improved the performance (thanks!). I doubt that the pivot to permutation conversion has more than a negligible time cost now?

The underlying jax.lax.linalg.lu_solve is accessible, should you want it.

What do you think?

[ahoenselaar]

I agree that we can close this issue.

…

On Tue, May 18, 2021 at 1:19 PM Peter Hawkins ***@***.***> wrote: I'm curious if you think this is still worth solving, given that #5880 <#5880> significantly improved the performance (thanks!). I doubt that the pivot to permutation conversion has more than a negligible time cost now? The underlying jax.lax.linalg.lu_solve is accessible, should you want it. What do you think? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#5826 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AA4VK57JLCLH7FAGILKMWNDTOLDUTANCNFSM4YDQ4C7A> .

[hawkinsp]

Thanks!