Why jax choose jaxpr instead of python ast ?

[edimetia3d]

Hi,

I am learning jax and some other python jit project, e.g. llvm-npcomp/numba, by now.

And I noticed that jax convert python source code to jaxpr as an IR by tracing the code execution hisotry, and at same time, llvm-npcomp/numba do the same thing by heavily using the python ast module.

For I am just a newbee to these jit techs, I couldn't see obivious benifits from "tracing the code", both "python ast" and "tracing" looks fine to me.

Was there some problem that only "tracing" solution could handle, while "python ast" not?

Sincerely.

[jakevdp]

This is an interesting question: I think the main reason is that the Python AST is far too general: JAX does not implement everything that Python implements: for example, transformable JAX programs are linear and functional: JAX variables cannot be modified in-place, and JAX does not have general control flow like if/else statements or loops! For that reason, the Python AST is not a great fit for representing the set of programs that JAX is able to transform and compile, and a more limited JAX-specific IR is a better choice. Does that answer your question?

[edimetia3d]

Thanks,

Your answer sloves a lot of questions in my mind. Using jaxpr to make things simple and clear is a good point.

There is one more little question. It's about the "Building jaxprs with tracing".

Assume that I want to use jaxpr as the jit IR in my new DSL, so I could use jax as my "jit compiler backend".

Because my DSL is simple and limited, it is easy to convert to jaxpr just by some text processing. Will I lose some core feature by "Building jaxpr with text processing" instead of "building jaxpr with tracing"? Should I also implement a tracing system in the new DSL?

[jakevdp]

I don't think there's anything special about tracing; if you just want to generate jaxprs directly from a parse tree of a custom DSL, that seems like a valid approach.