Add JAX glue for tvm-ffi

[danieldk]

[varshith15]

hey @danieldk can i pick this up? did some work here need to test and cleanup a bit but the idea is to add jax support for tvm-ffi using jax-tvm-ffi to register kernel calls as XLA custom calls so they run in jax.jit unlike eagar pytorch
let me know what you think

[varshith15]

hey @danieldk any chance you looked at my code? thoughts?

[danieldk]

Hi @varshith15! Sorry, I hadn't seen your comments. We are working on JAX support using jax-fvm-ffi as well. We probably don't want JAX support code in the kernels Python package itself (outside kernel layers, we also do not have a lot of Torch-specific code, kernels is mostly for fetching + loading), but we are still trying some different approaches.