Why do we need the Trace class besides the MainTrace class?

[eaplatanios]

I have been going over the Autodidax tutorial (thank you for this amazing learning resource!) and I can’t quite tell what the need for the Trace class is. Couldn’t we just have pure, lift, and process_primitive defined in MainTrace and then just pass around references to a single shared interpreter instance? Is there a particular need to carry additional state in Trace class instances or something like that? Thank you!

[patrick-kidger]

I think this might be a design choice based around higher-order primitives. In this case the e.g. JVP implementation would like to "re-dispatch" and itself call a JVP. The MainTrace is held in common for the whole invocation, but the surrounding JVP and the "sub JVP" have separate Traces.

(I'm willing to be wrong on the above, if someone better-informed knows otherwise.)

One can definitely imagine other tracing models than JAX's, where e.g. a single Trace object is passed around, or where things operate based purely on data dependencies rather than dynamic scopes.

[eaplatanios]

Thanks for the prompt response! I guess the main part I don’t understand is that throughout that tutorial, Traces appear to be wrappers over MainTrace instances that hold no additional state. JVPTrace for example is constructed by simply wrapping a MainTrace and doesn’t hold anything else internally. Looking at this again it appears that we never subclass MainTrace and so I wonder if the current design is equivalent to one where MainTrace is abstract and subclassed to e.g., create a single shared JVPMainTrace for a given stack level, avoiding the need for multiple trace instances per MainTrace. Does that sound plausible?