Does it make sense to use Cython with JAX?

[mario-linov]

Hi.
I am thinking on starting to use JAX. I was wondering whether it makes sense to use Cython and the JIT compilation of JAX.
Maybe they do not add execution speed to each other.
Thanks.

[jakevdp]

Thanks for the question!

Cython and JAX JIT are mostly incompatible: The goal of Cython is to facilitate writing of CPython extensions to call from Python and execute on the CPU. JAX, on the other hand, provides a domain specific language for array computing that lowers to XLA, for compiled execution on CPU, GPU, TPU, and other architectures. The two don't naturally intersect, because CPython and XLA are entirely different backends. That said, it might be useful in some cases to use Cython to implement functionality that could be called from a JAX program running on CPU, when functionality cannot be easily expressed in XLA. We're currently working on adding better ways to interface with such foreign functions; see #12632.

Does that answer your question?

[mario-linov]

Thank you very much for your explanation. Clear and complete 👌.