Closed adam-hartshorne closed 1 week ago
FYI, the semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated in the latest JAX release, and looking at the torch2jax is used.
https://github.com/jax-ml/jax/releases/tag/jax-v0.4.35
Thanks for the update! This is probably now the last wake-up call, I'll migrate to the new FFI interface ASAP.
Thankfully, the register_custom_call_target still works in tests, so no breaking changes yet.
FYI, the semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated in the latest JAX release, and looking at the torch2jax is used.
https://github.com/jax-ml/jax/releases/tag/jax-v0.4.35