SIGN IN SIGN UP

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