Explore using the new JAX FFI instead of custom call based on pybind11 #132
Open
Description
JAX recently added a proper FFI, which removes the need to create Python bindings for the functions we want to call (removing the need for something like #107).
See the documentation here: https://jax.readthedocs.io/en/latest/ffi.html
I'm not sure what the consequences for pre-built sphericart-jax
wheels are. This will depend on how stable the ABI used for FFI is. We might have to do the same trick as with PyTorch versions (having multiple copies of the sphericart-jax shared library, each compiled against a different JAX version).
Metadata
Assignees
Labels
No labels