Jax-Transform based Sow Implementation

[samanklesaria]

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).

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 Variables 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.

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:

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:

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:

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:

@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:

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:

_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 binds 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:

@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 reachessow[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 jits 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
test_sow.py