Migrate `equinox` to using `jax.extend`?

Open michael-0brien opened this issue 4 months ago • 1 comments

As per JAX's extension guides, jax.extend is an API for calls into the JAX internals. For getting deprecation issues resolved in a robust way like the recent deprecations in jax.interpreters in JAX 0.7.1 (#1081), should we make the appropriate feature requests in jax.extend? I expect that JAX would mutually benefit from this as equinox is a big use-case of an extension in the JAX ecosystem! This will presumably be a cumbersome endeavor so the developers should feel free to delegate tasks :)

Aug 21 '25 18:08 michael-0brien

I honestly don't have super strong feelings. jax.extend doesn't actually come with any greater stability guarantees than e.g. jax._src, after all! :)

Aug 22 '25 08:08 patrick-kidger