Calling Python JAX functions from a custom C++ XLA op

[denizokt]

Hi all, I want to see a use-case like the following is possible: I have a function defined in JAX Python along with its gradient (through autodiff). I want to implement a fast operation involving this function (one example is integration) that would need to use this JAX function and its gradient as a black-box.

I would like this operation to be as fast as possible, so I would like to implement a custom OP in CUDA/C++ to do it, but I would like to be able to call the JAX function from within this operation.

Is this possible at all? I am willing to do some digging and plumbing to get it to work if needed, but could someone point me in the right direction?

Thank you!
Deniz