Jax-Transform based Sow Implementation

# Proposal: a functional `sow` transform for extracting intermediates

## Summary

Replace the Module-bound intermediate-capture system (`self.sow(...)` + `nnx.capture` ) with a **free function** `sow` and
a **function transform** `capture`, built on a small JAX primitive and a jaxpr interpreter (the same idea as oryx's `harvest`).

```python
def f(x):
y = layer(x)
sow(y, name="features") # tag any value, anywhere
return head(y)

out, collected = capture(f)(x) # collected["features"] == (y,)
```

No module, no variable collection. `sow` is a plain call you can drop into *any* function — a layer, a loss, a bare arithmetic expression — and `capture` harvests every tag by name.

## Motivation

Today, capturing an intermediate requires:

- the value to live inside an `nnx.Module.__call__`
- a `self.sow(nnx.Intermediate, key, value)` call keyed by module path,
- wrapping the call site in `nnx.capture(...)`

This couples "I want to look at this array" to the module/variable system. You can't tag a value in a plain function, in a loss, or in a gradient. The keying is by module structure rather than by an explicit name you choose.

The functional version decouples tagging from modules entirely. It's ~200 lines, composes with `jit`/`scan`/`cond`, and extends more cleanly to gradient capture than the current approach.

### Why this matters now: the Hijax transition

The larger motivation is that we want to move flax `Variable`s onto **Hijax**, so that variable state becomes real mutable references, enabling mutation of values captured in closures, among other things. The current intermediate-capture
system works by mutating a variable collection during the forward pass in a way that changes the static type of the collection. But **QDD for Hijax is deprecated**, making a pure Hijax transition incompatible with current implementation of `sow`.

The transform proposed here sidesteps this entirely: it captures intermediates by **reading a jaxpr**, not by mutating state. There is no variable collection to accumulate into and no QDD dependency — `capture` can be a standard jax transform which recovers the tagged values by interpretation. This is therefore a `sow` that simplifies the move to Hijax.

## How it works

### 1. A `sow` primitive that is an identity

`sow` is a JAX primitive whose implementation, lowering, and every transform rule are the identity — it does nothing to the value. It exists only to leave a named marker in the jaxpr.

```python
sow_p = core.Primitive("sow")
sow_p.multiple_results = True
sow_p.def_impl(lambda *xs, **__: list(xs))
mlir.register_lowering(sow_p, lambda ctx, *args, **__: args) # identity
ad.deflinear2(sow_p, lambda cts, *_, **__: cts) # identity under autodiff
batching.primitive_batchers[sow_p] = lambda args, dims, **p: (sow_p.bind(*args, **p), dims)

def sow(value, *, name):
leaves, treedef = tree_flatten(value)
out = sow_p.bind(*leaves, name=name, tree=treedef) # name + treedef ride along as params
return tree_unflatten(treedef, out)
```

Because it is a true identity, a `sow`-containing function still runs correctly under plain `jax.jit` — the markers are simply discarded.

### 2. An effect so bare sows survive DCE

A bare `sow(loss, name="loss")` whose result is unused would normally be dead-code eliminated. Attaching a JAX effect keeps it in the jaxpr:

```python
class SowEffect(effects.Effect): pass
sow_effect = SowEffect()
effects.lowerable_effects.add_type(SowEffect) # + control_flow / remat allow-lists
sow_p.def_effectful_abstract_eval(lambda *avals, **__: (list(avals), {sow_effect}))
```

### 3. `capture`: trace to a jaxpr, then interpret it

`capture(f)` stages `f` to a jaxpr and walks it. Every equation is replayed faithfully **except** `sow_p`, which records its (name → pytree) into a dict and passes the value through:

```python
def capture(f):
def wrapped(*args, **kwargs):
cj, out_shape = jax.make_jaxpr(f, return_shape=True)(*args, **kwargs)
out_flat, collected = _run(cj.jaxpr, cj.consts, *tree_flatten((args, kwargs))[0])
out = tree_unflatten(tree_structure(out_shape), out_flat)
return out, {name: tuple(vals) for name, vals in collected.items()}
return wrapped
```

The interpreter is a standard eval loop over the jaxpr:

```python
def _run(jaxpr, consts, *args):
env, collected = {}, {}
read = lambda v: v.val if isinstance(v, core.Literal) else env[v]
env.update(zip(jaxpr.constvars, consts))
env.update(zip(jaxpr.invars, args))
for eqn in jaxpr.eqns:
invals = [read(v) for v in eqn.invars]
if eqn.primitive is sow_p:
pytree = tree_unflatten(eqn.params["tree"], invals)
collected.setdefault(eqn.params["name"], []).append(pytree)
outs = invals # pass through
else:
outs = ... # bind the primitive as usual
env.update(zip(eqn.outvars, outs))
return [read(v) for v in jaxpr.outvars], collected
```

`collected` maps each name to a **tuple of pytrees in call order**, so sowing the same name twice gives you both.

## 4. `perturb`: capturing gradients

The same machinery captures **cotangents** with a one-liner. `perturb` is an identity whose `custom_vjp` backward rule sows the incoming gradient. Because the backward rule is staged into the gradient jaxpr at trace time, `sown(jax.grad(f))` harvests it:

```python
@partial(jax.custom_vjp, nondiff_argnums=(1,))
def _perturb(value, name):
return value
_perturb.defvjp(lambda value, name: (value, None),
lambda name, _res, g: (sow(g, name=name),)) # sow the cotangent

def perturb(value, *, name):
return _perturb(value, name)
```

Usage:

```python
def loss(x):
return jnp.sum(perturb(x, name="grad_x") ** 2)

grad, collected = sown(jax.grad(loss))(x)
# collected["grad_x"][0] == d(loss)/d(x)
```

This is something the current Module-bound system handles far less elegantly.

### 5. Control flow

In our interpreter, sows nested inside higher-order primitives are handled by recursing into their sub-jaxprs:

- **`jit` / `closed_call` / `custom_jvp` / `custom_vjp`** — inline the sub-jaxpr and collect directly (see below).
- **`scan`** — a sow fires once per iteration; the interpreter turns the sown values into extra `ys`, so `scan` stacks them along the scan axis. You get one pytree per name with a leading length dimension.
- **`cond`** — sown values become extra branch outputs. JAX's own "all branches must match" check enforces that every branch sows the same names and shapes.

#### Inlining the "call" primitives

`jit`, `closed_call`, and the two `custom_*` primitives are all *call* primitives: each one wraps an entire sub-jaxpr as a parameter and, when bound, runs that sub-jaxpr as an opaque unit. A sow living inside that sub-jaxpr is therefore invisible to us if we just `bind` the primitive — its markers stay sealed behind the call boundary.

The fix is to not bind the primitive at all. Each of these primitives keeps its body under a known parameter key, so we look the key up, pull out the sub-jaxpr, and hand it to the *same* interpreter:

```python
_CALL_PRIMS = { # primitive -> param key holding its sub-jaxpr
pjit.jit_p: "jaxpr",
core.closed_call_p: "call_jaxpr",
custom_jvp_call_p: "call_jaxpr",
custom_vjp_call_p: "call_jaxpr",
}

elif prim in _CALL_PRIMS:
sub = eqn.params[_CALL_PRIMS[prim]]
jx, cs = (sub.jaxpr, sub.consts) if isinstance(sub, core.ClosedJaxpr) else (sub, [])
outs, sub_coll = _run(jx, cs, *invals) # recurse with the SAME interpreter
_merge(collected, sub_coll) # fold child names into ours, in call order
```

Three things fall out of this:

- **Correctness is unchanged.** `_run` re-emits the sub-jaxpr's equations one for one and returns its outputs, so the values flowing out are exactly what `bind`-ing the primitive would have produced. The only added behavior is that the recursion also returns a `collected` dict, which `_merge` folds into the parent's — extending each name's call-order list rather than overwriting it, so sows at different nesting depths accumulate into one flat namespace.
- **Nesting is free.** Because we recurse with the *same* `_run`, a sow buried under `jit(jit(...))` or inside a `custom_vjp` that itself contains a `scan` is reached with no extra cases — each layer just inlines the next.
- **The call boundary is dissolved during harvest.** We never re-`bind` `jit_p`, so an inner `jit` is effectively unrolled into the surrounding interpretation. This is a trace-time flattening only: if the outer `capture(f)` is itself jitted, everything re-fuses into a single XLA computation, so there's no runtime cost — you simply lose the inner compilation cache boundary while harvesting.

#### Why a sow in a *backward* rule needs no special handling

The `custom_jvp` / `custom_vjp` entries in `_CALL_PRIMS` point at each primitive's **primal** `call_jaxpr` — the forward computation. That seems to leave a gap: what about a sow placed in a `custom_vjp`'s *backward* rule, like `perturb` uses? It turns out there is nothing extra to do, and the reason is a matter of *ordering*.

The key fact is that **autodiff is a trace-time transformation, not a runtime one.** When you write `capture(jax.grad(loss))`, `capture` calls `jax.make_jaxpr(jax.grad(loss))`. Building that jaxpr forces JAX to actually *perform* the differentiation while tracing: it traces the forward pass, then traces the backward rules to stage the gradient computation. A `custom_vjp`'s backward rule is just an ordinary Python function, and tracing it runs its body line by line — so a `sow(...)` call sitting inside it `bind`s a `sow_p` equation into the jaxpr being built, exactly like a sow anywhere else.

Crucially, that backward rule is *not* stored back inside the `custom_vjp_call_p` equation. The `custom_vjp_call_p` primitive that survives into the jaxpr represents only the **forward** call; the backward equations (including the sow) are spliced into the enclosing jaxpr as flat, top-level equations. So by the time `capture` walks the jaxpr, the backward sow is no longer "inside" a custom_vjp at all — it is a plain `sow_p` equation that the ordinary `if prim is sow_p` branch of `_run` picks up directly.

##### Concrete example

Take `perturb`'s pattern — a `custom_vjp` whose backward rule sows the incoming
cotangent:

```python
@jax.custom_vjp
def h(x):
return x ** 2
def h_fwd(x):
return h(x), x
def h_bwd(x, g):
gx = 2 * x * g
sow(gx, name="grad_h") # sow lives in the BACKWARD rule
return (gx,)
h.defvjp(h_fwd, h_bwd)

def loss(x):
return jnp.sum(h(x))
```

`jax.make_jaxpr(jax.grad(loss))(jnp.arange(3.))` produces:

```
{ lambda ; a:f32[3]. let
b:f32[3] = custom_vjp_call[ # <-- FORWARD call only; call_jaxpr is just x**2
call_jaxpr={ lambda ; c:f32[3]. let d:f32[3] = integer_pow[y=2] c in (d,) }
bwd=h_bwd fwd=h_fwd name=h
] a
_:f32[] = reduce_sum[axes=(0,)] b
e:f32[3] = broadcast_in_dim 1.0
f:f32[3] = mul 2.0 a # <-- these three equations ARE the backward pass,
g:f32[3] = mul f e # staged flat into the top-level jaxpr
_:f32[3] = sow[name=grad_h tree=*] g # <-- the backward sow, now an ordinary sow_p equation
in (g,) }
```

Notice what did and didn't happen:

- The `custom_vjp_call` equation's `call_jaxpr` contains only `integer_pow`: the forward `x**2`. There is no sow inside it. That is why `_CALL_PRIMS` inlining the *primal* jaxpr is sufficient: the primal genuinely holds no sow.
- The backward multiplies and the `sow[name=grad_h]` equation sit at the **top level** of the gradient jaxpr, siblings of the forward call.

So when `_run` walks this jaxpr it inlines the primal `custom_vjp_call` (finding nothing to collect), then reaches`sow[name=grad_h]` as a first-class equation and records `g` under `"grad_h"` — no custom_vjp-aware logic required.

`while_loop` is intentionally unsupported (a dynamic trip count can't produce a fixed-size collection); `remat` is likewise out of scope.

## Limitations

- `while_loop`, `remat` are unsupported.
- `sown` re-traces `f` to a jaxpr (one extra trace); harvested runs are
interpreted in Python rather than staged (though nested `jit`s still lower
normally).
- Names are a flat namespace; collisions accumulate into the call-order tuple.

## Reference implementation

A complete, tested ~200-line implementation lives in `sow.py`, with tests in
`test_sow.py` covering top-level, `jit`, `scan`, `cond`, `custom_vjp` backward
capture, and the unsupported-`while_loop` error path.

[sow.py](https://github.com/user-attachments/files/30474983/sow.py)
[test_sow.py](https://github.com/user-attachments/files/30474982/test_sow.py)

Start by locating the existing Module-bound sow/capture implementation and the proposed entry points: sow, capture, _run, and perturb. Read the surrounding JAX transform and jaxpr-interpreter code, then trace the proposal's jit, scan, cond, autodiff, and nested-call cases. Done means named pytrees and gradients are captured without variable collections across the listed transformations.

Written by the indexing model from the issue text.