Interpreter idioms - staging to Jaxpr, vs. writing own tracers?

[femtomc]

Hi all!

I've written a few interpreters using JAX now, and typically I follow the following pattern:

1. Stage a program, getting a ClosedJaxpr out.
2. Initialize an environment and an interpreter which walks the ClosedJaxpr, processing each primitive in turn.
3. Return the values for the final outvars - or other quantities computed during the interpretation process.

I've never had to write my own Tracer type -- but I'm curious how this approach compares to writing your own Tracer type, etc?

Is there a downside to staging first?

[jakevdp]

Others might chime in, but in general the easiest approach to implement a custom transformation is using what we often call initial style: staging out to a jaxpr and using a jaxpr interpreter. The alternative is final style, implementing the transformation on-the-fly using a custom tracer. This is harder, but the advantage is often easier debuggability because the transformation happens within the context of the traced code, rather than as a post-processing step after it is traced. It also is more flexible, e.g. it allows for Python control flow to depend on traced data in some cases. But final-style transformations are somewhat more difficult to implement.

If you want to see a side-by-side comparison of the two approaches, one place that exists in the JAX codebase is in the two implementations of the sparsify transform; you can see the code here: https://github.com/google/jax/blob/main/jax/experimental/sparse/transform.py#L456-L459

There's also some discussion of these two approaches in autodidax.

[femtomc]

Thanks Jake!