What does pure, lift, and sublift mean in JAX core?

[leiteg]

After reading the JAX core from scratch and reading some parts of the JAX source code, I was confused as to what pure, lift and sublift methods in core.Trace mean? Do they simply take a JAX array and create an ShapedArray from them?

[mattjj]

Thanks for the question!

It may be tricky to answer in full detail, so instead I hope this summary answer is still useful:

• trace.pure is called on non-Tracer values, e.g. numpy.ndarrays or Python builtin number type instances, and produces a Tracer for the trace meant to represent having trivial context added (e.g. a zero tangent for JVPTrace, a None batch dim for BatchTrace, etc);
• trace.lift is similar but is always called on Tracer values, specifically ones that don't belong to the trace (e.g. if we're doing jvp-of-jvp, or vmap-of-jvp, then the inner one's trace.lift might be applied to a Tracer from the outer one), and also produces a Tracer for the trace meant to represent trivial context;
• trace.sublift is similar but is always called on Tracer values, specifically ones that do belong to the same trace, but from outside the scope of special kinds of higher-order primitives (like custom_jvp).

The last one has no analogue in Autodidax (the JAX core from scratch tutorial).

Does that help?

[leiteg]

Thanks @mattjj, that helps a lot! So does that mean that the method full_raise of Trace simply takes any value and delegates it to one of pure, lift, and sublift?

I need some time to process this information, but I am satisfied for now, marking as answered 😃