Using cupy within jax primitives

[sunilkpai]

Hi, I am currently working on a pull request to use cupy's cuSPARSE wrappers to implement preconditioned iterative solvers. This would be very useful to accelerate inverse designs of large 3D photonic structures in a scaled up implementation of ceviche.

I was wondering if there are any issues running cupy within a custom primitive, essentially the same as what ceviche provides in autograd, except using cupy and jax instead. Admittedly, I have not played with jax's interface yet and I am also relatively new to autograd.

I am also potentially interested in implementing some of these algorithms in jax along with the cg solver. Let me know if this is a feature you would like and I can get started on a PR (I will work on it in my free time). In that case, I wouldn't require using cupy-jax interoperability.

[hawkinsp]

I don't think anyone has tried this! It may well work, but the main caveat is that you won't be able to jit the result (not without doing more work). Let us know what happens?