Stable API: Initial Style vs Final Style transforms

[erick-xanadu]

Hi,

we are thinking of writing some JAX transforms and we have been debating about whether to use Initial Style vs Final Style. We have noticed that the classes for Trace and Tracer are not in the stable API as defined in this list. This means that Final Style transforms cannot be implemented using the stable API. It also looks like one can implement Initial Style transformations using only stable API / python.

How unstable are the classes needed to implement a Final Style transform (like Trace and Tracer)?

EDIT: I know technically a lot of Trace and Tracer is left to be implemented in child classes, but should one put consideration into the fact that they are not in the stable API?

Thanks!

[jakevdp]

Thanks for the question – I think @mattjj and @froystig might also want to weigh in here, but I think it is not our intent that the tools for writing custom transformations would be part of JAX's public/stable API. Things at that level are still considered internal APIs, and we don't expect most users to have to interact with them. We might think about putting pieces of it into jax.extend (see https://jax.readthedocs.io/en/latest/jep/15856-jex.html) but that is still evolving at the moment.

Out of curiosity, what is the use-case that you have in mind?

[froystig]

Nothing here is public nor stable. If you're digging in anyway then, if at all possible, write a Jaxpr interpreter that begins by calling jax.make_jaxpr.

Parts of the Jaxpr IR are likely eventual candidates for jax.extend. Traces, tracers, and others things to do with "final style" are all much less so. And note that jax.extend is never stable.

[erick-xanadu]

Thanks @jakevdp and @froystig !

Out of curiosity, what is the use-case that you have in mind?

We are working on creating a JIT compiler for a data science library. The data science library works a little bit like an interpreter. We are exploring using custom primitives with concrete evaluations that would allow this interpreted library to be easily implemented in Python. At the same time, some transformations that would allow the program to be optimized and lowered into MLIR via the MLIR interpreter and fed into our custom compiler.

If you have any questions of comments, happy to chat! Thanks for all the help :)

[jakevdp]

Thanks for the context! It sounds like a cool project. That said, this use-case sounds a bit out-of-scope for what JAX can provide as far as stable APIs for downstream consumption. I wonder if you might not be better served by developing and maintaining your own tracing stack that's inspired by JAX's approach? You could follow autodidax to get the basic setup and then work from there.