(Interpreter idioms for lax primitives like lax.cond) Apply interpreters across branches - what happens to literals in branches?

[femtomc]

Hi all!

I'm working on a transformation which uses an interpreter and a set of rules to parametrize the functionality on primitives.

Working with straight line programs, I've run into no problems with my interpreter. However, I'm trying to extend the transformation to operate on jax.lax.cond - which holds a reference via branches to the two ClosedJaxpr branch instances.

The right thing I think here is to apply the interpreter to each branch (keeping the return values consistent, as expected) -- and then return a value again using jax.lax.cond with the transformed branches.

However, I'm running into issues for certain branch functions -- like the following:

{ lambda ; a:f32[]. let b:f32[] = div a 2.0 in (b,) }
{ lambda ; a:f32[]. let  in (0.0,) }

Imagine I'm writing a forward mode AD interpreter which lifts values to Dual numbers. I can apply the transform to the first branch, and, to be consistent -- my transform will lift all constants (like 2.0) to Dual.new(2.0, 0.0) (for example).

However, when I try and inspect the closed_jaxpr.literals for each branch in jax.lax.cond -- I get the following printout:

{ lambda ; a:f32[]. let b:f32[] = div a 2.0 in (b,) }
# print(branch_1.literals)
[]
{ lambda ; a:f32[]. let  in (0.0,) }
# print(branch_2.literals)
[]

This seems weird to me because 2.0 is a literal, and so is 0.0 -- and then when I apply the transform and use the result in a new cond call, it of course will complain about non-matching Pytree types across branches.

Any ideas why this occurs?

[mattjj]

Sounds like you're cooking up something cool!

IIUC, you're asking why 0.0 and 2.0 don't show up in closed_jaxpr.literals, right?

If so, it's just a result of confusing naming choices. We named the attribute on ClosedJaxpr "literals" because:

1. we were tired of using "consts" everywhere, and
2. we wanted to underscore that that field could never contain Tracers, and had to contain raw array-like stuff.

(Speaking of confusing names, ClosedJaxpr should really be called PartiallyAppliedJaxpr...)

But actual literals, as in constants whose values appear directly in the program data structure, don't go in in that field at all. They go somewhere else.

There's another kind of object called core.Literal, which go in core.JaxprEqn.invars and also in core.Jaxpr.outvars (which should be called in_atoms and out_atoms respectively, since they are lists of core.Atom = Union[core.Var, core.Literal], not lists of core.Var). That's where the 2.0 and 0.0 are. For example, you should be able to see the Literal(0.0) object if you do branch_2.jaxpr.outvars[0].

Does that make sense? Let me know if I'm misunderstanding!

[femtomc]

Thanks @mattjj — indeed, you understood correctly.