axlearn icon indicating copy to clipboard operation
axlearn copied to clipboard

Published 20 hours ago • verified publisher icon apple

Remove usage of jax.core and jax.interpreters

Open samos123 opened this issue 11 months ago • 1 comments

Those were internal APIs and will be deprecated / removed.

jax.core.Primitive should use jax.extend.core.Primitive

There is also an isinstance check for jax.core.Tracer which may require some more thought on how to replace the logic.

I can help with this work in collaboration with Jax team.

samos123 avatar Dec 09 '24 18:12 samos123

Relevant info:

  • https://jax.readthedocs.io/en/latest/api_compatibility.html
  • https://jax.readthedocs.io/en/latest/jax.extend.html

jakevdp avatar Dec 09 '24 18:12 jakevdp