Defer bind with unboxed values to next interpreter

[femtomc]

How do I do this?

There's this section in Autodidax which roughly covers what I was thinking about. The issue is (unlike make_jaxpr), my inner interpreters might call bind when process_primitive.

I have primitives that are active for one interpreter, but if I perform another transformation first - that other transformation might call bind with unboxed values (that transformation creates a new jc.new_main(MyTrace) as main).

It seems like this removes the possibility of my outer transformation processing that primitive.

Edit: is full_lower what I want? Conceptually, I want my process_primitive to call bind with the Tracer values corresponding to the next interpreter on the stack (or, just the values -- to default to the evaluation interpreter, or whatever is on the bottom of the stack).

[femtomc]

I figured out how to do this - but I had to rewrite a few of my transformations to always use my own interpreter pattern (which always stages out to a Jaxpr, before walking -- so each transformation which uses this interpreter gets to see every primitive).

I'm not sure there's a general way to do it without doing this because (as mentioned in Autodidax - JAX has dataflow assumptions as an optimization, so I'm sort of breaking those assumptions so I had to go lower level and make my own interpreter stack -- this is totally okay for my use cases, however).