Add `exclude_argnames` on program_order.
**Why?** program_order(enforce=True) traverses the jaxpr and adds opt_barrier of the form `opt_barrier((prev_outs, cur_inputs))`. It does not recurse into HOPs though. One consequence of adding opt_barriers (other than disallowing all kinds of optimizations) is that it also alters the lifetime of some arrays (I'll give an example below). So what `exclude_argnames` on `program_order(enforce=False)` does is it allows us to skip some `cur_inps` from the `opt_barrier((prev_outs, cur_inps))`. This opt_barrier is inserted by an outer program_order(enforce=True), so `enforce=False` is nested. Note that exclude_argnames only works when `enforce=False`. Also, if some arrays are closed over, then those can't be excluded. To exclude, arrays should be inputs and added to `exclude_argnames` (maybe we can support closed over inputs somehow but not today). To make this work, we added a `program_order` primitive but it's very lightweight since it piggybacks on `eval_jaxpr_p`! **Example of opt_barrier altering lifetimes** ``` x, y = opt_barrier((x, y) # Will force both x and y to be ready here! f(x) g(y) ``` Without the opt_barrier: ``` f(x) # x can be ready here g(y) # y can be ready here ``` PiperOrigin-RevId: 972804273
Y
Yash Katariya committed
94b575cd8259c9c86ba98cb0eb053226b8aad06b
Parent: d551c34
Committed by jax authors <google-ml-automation@google.com>
on 8/28/2026, 9:55:53 PM