Open f0uriest opened 6 months ago
Definitely something we'd be interested in including!
On licensing issues: if you've used the original and then modified it, then we can probably include it whilst reproducing their own license in the implementation file, and in our own top-level LICENSE.
On tests: pytrees definitely need testing. You should also be able to register it with our jvp/vmap/etc. tests that already exist.
Hi @patrick-kidger , I've really enjoyed playing around with Lineax!
You have several iterative methods for square systems (CG, GMRES, BICGSTAB) but nothing for least squares problems.
I have a JAX port of
scipy.sparse.linalg.lsmr
that I should be able to get working with lineax, but wanted to check first: