Aaditya Chandrasekhar

@amir90 I am facing the same issue. Were you able to find any resolution to this?

I am trying to use [PyAMG](https://pyamg.readthedocs.io/en/latest/generated/pyamg.krylov.html) for solving my PDEs. I thought I could wrap the solver inside `jax.lax.custom_linear_solve` to ensure that I can also perform the adjoint solve. However,...

Thank you so much @shoyer ! This worked For folks from the future who might be running into similar issue [here](https://gist.github.com/aadityacs/3c5be1210311bdb9ba946f42df767a73) is the implementation I may still be missing something...

I think we should be able to leverage already existing solvers on GPU(s) to speed up the computation. For instance see this [notebook](https://colab.research.google.com/drive/1aHF6k1CNgmCel64er6b9UD2il9RdF0IM?usp=sharing) where I am simply calling `cupy` code...

Any updates on this? Thank you!