New solvers

[patrick-kidger]

• Stabilised Thomas algorithm. (For tridiagonal solves.)
• Upper Hessenberg. (Useful inside GMRES? Needs to offer a pseudoinverse solution.)
• Incremental GMRES
• Handle "diagonal + low-rank" operators using Woodbury.

[joglekara]

Thanks for the library. As they said, and another one...

Do you have any intuition whether it would be useful to swap in the cusparse kernels for, say, a tridiagonal solve, and add a custom_linear_solve for the gradient or is that more trouble than worth the performance boost (if any!)?

At the moment, I have a custom written broadcast-based function similar to the one in here but with a, b, c, d being 2D. I suppose I could've used vmap but I thought to trust broadcasting.

But would something like this be better than what XLA comes up with?

https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-t-gtsv2stridedbatch

Thanks again

[patrick-kidger]

I think that would definitely be useful. Note that we actually already have the gradient via our custom primitive (linear_solve_p) so the only difficulty with using cusparse might be adding batch dimensions; off the top of my head I'm not sure if they support that.

But at least in the unbatched case, I reckon that probably would give a performance boost. (Plus I believe their tridiagonal algorithm is stabilised, whereas our current Thomas algorithm is unstablised.)

The only reason we haven't pursued this is that it's not straightforward to make this happen in JAX -- calling out to custom kernels like this isn't very well supported.

[joglekara]

Couple of disparate thoughts

• They do already have a separate call that supports the batched dimension
• Sounds like I should port over to your solver then
• Is something like this what is needed? Add lax.tridiagonal_solve(..), lowering to cusparse<T>gtsv2() on GPU. google/jax#6849

[patrick-kidger]

Something like that!
@packquickly may recall why we didn't use the built-in tridiagonal solve algorithm. IIRC it doesn't support batching, and that has to be patched in? (E.g. by vmap'ing jax._src.lax.linalg_tridiagonal_solve_jax.)

[quattro]

Would there be any interest in randomized solvers (e.g., via sketching or randomized trace estimation)?

[patrick-kidger]

Absolutely!