pipekit under JAX

Summary¶

Decision. JAX support stays in two places. pipekit core stays JAX-free and carrier-agnostic, but refuses or flags the operators that are wrong under tracing. pipekit-jax holds everything that imports JAX or Equinox: today’s model wrapper, forward model and adjoint helpers, and, when a project needs gradients through operator parameters, the operator-as-PyTree layer the scoping report designed. Five workstreams.

• A. Trace guards in pipekit core. Cache keys every tracer of one shape to the same entry, so a pipeline that caches twice returns the first result for the second call (A1). Snapshot, Sink, Profile and Tap run once, at trace time, on a tracer; Snapshot leaks it (A2).

• B. Operators as PyTrees in pipekit-jax. A JaxOperator(Operator, eqx.Module) mixin and a JaxSequential whose operators are a tuple field let gradients reach parameters and let eqx.tree_at replace a step (B1). Their configs hold arrays, which breaks state as JSON; the config needs an array policy (B2).

• C. JAX-native control and observation. A CondBranch and CondSwitch on jax.lax.cond and jax.lax.switch (C1), and a DebugTap on jax.debug.callback, which runs on every call rather than once per trace (C2).

• D. Transformation wrappers. VMapOver, Checkpoint and ScanSequential wait for a project that needs them; truncated_scan already covers the cycle layer’s scans.

• E. Docs. One page on build time against trace time (E1), and fixes to pipekit-jax’s README and install notes (E2).

End state. A pipeline of JAX operators runs under jit, vmap and grad as it does today. An operator that cannot be traced says so instead of returning a wrong value, a pipeline whose parameters are learned is built from PyTree operators, and its trained weights round-trip through the registry the way JaxModelOp’s do.

Ownership rules¶

pipekit core owns composition and the operators that misbehave under tracing; pipekit-jax owns everything that imports JAX (Figure 1).

1. pipekit core owns Operator, Sequential, Graph, Fanout and the control, observation and caching operators. It never imports JAX. Its guards (A) detect a tracer through sys.modules, only when JAX is already imported.

2. pipekit-jax owns every class that subclasses or wraps an eqx.Module, every operator built on a jax.lax primitive, and the JAX-side adjoint helpers. It depends on pipekit, pipekit-cycle (which Report 16’s B moves to the [diffrax] extra), Equinox and JAX.

Neighbours (pipekit-train, pipekit-cycle, the algorithm libraries) follow the stack rules; for this page the one that matters is that pipekit-train differentiates the module a JaxModelOp wraps, never a pipekit Sequential (stage 2).

What ships¶

Table 1 lists what pipekit-jax 0.0.2 exports beside the scoping report’s names; the italic rows were planned and never built.

The running example: one retrieval under jit, vmap and grad¶

The sensor is the five-band sensor of Report 17, unchanged: $y_b = e^{-k_b \alpha} + \varepsilon_b$ with $k = (0.002, \dots, 0.010)$, $\sigma = 10^{-3}$ and $\alpha \sim \mathcal{U}(0, 300)$. The retrieval inverts each band, $z_b = -\log y_b / k_b$, and averages the five with inverse-variance weights. To first order $\operatorname{var} z_b = \sigma^2 / (k_b y_b)^2$, which gives (1):

The test set is 500 scenes; stage 2 adds 2,000 labelled scenes, as from controlled releases.

Stage 1 · jit, vmap and grad through a pipeline (pipekit → JAX)¶

TL;DR. Run the retrieval as a plain pipekit Sequential under the three transformations, and use the gradient as each scene’s error bar.

Problem. A pipeline traces if every _apply is pure JAX and dispatch does not intercept the tracer. Operator.__call__ builds a graph node only when an argument is a Node (operator.py:160), and a JAX tracer is not one. The noise propagates through the retrieval as

so if grad is right ((2)), the errors divided by $\sigma_{\hat\alpha}$ have an RMS of 1.

1
2
3
4
5
6
7
8
9
10
11
12
# Shapes: n = 500 scenes, B = 5 bands
# z_b = −log y_b / k_b
invert = Lambda(lambda y: -jnp.log(y) / K)              # (B,) → (B,)   today
# α̂ = Σ_b w_b z_b / Σ_b w_b,  w_b = (k_b e^{−k_b z_b})²
combine = Lambda(weighted_mean)                         # (B,) → ()     today
retrieval = Sequential([invert, combine])                               # today
jax.jit(retrieval)(y[0])          # → 215.44 (α = 215.29); = eager      today
jax.make_jaxpr(retrieval)(y[0])   # → one jaxpr, 12 equations           today
jax.vmap(retrieval)(y)            # (n, B) → (n,): RMSE 0.262           today
# σ_α̂ = σ ‖∇_y α̂‖
s = SIG * jnp.linalg.norm(jax.vmap(jax.grad(retrieval))(y), axis=1)  # today
# RMS of σ_α̂ 0.258; RMS of (α̂ − α) / σ_α̂ = 0.97

All three transformations pass through the pipeline unchanged. The jitted output agrees with the eager one to 10⁻¹², because it is the same arithmetic in one jaxpr. The gradient is right too: the linearised error bar has an RMS of 0.258 against an actual RMSE of 0.262, and the standardised errors have an RMS of 0.97, within the 0.03 that 500 scenes allow. None of this needed pipekit-jax; it needs only operators whose _apply is pure JAX. Graph with a diamond (one node feeding two) and Fanout behave the same under jit and vmap (stage 3).

Stage 2 · Gradients with respect to an operator’s parameters (pipekit-jax → Optax)¶

TL;DR. Learn the absorption coefficients from labelled scenes, and see that the gradient cannot reach them through the pipeline.

Problem. The coefficients are known only to 10 %, as from an old line list. With 2,000 scenes whose $\alpha$ is known, fit $\log k$ by minimising

where the inversion step holds $\log k$ as an eqx.Module wrapped in a JaxModelOp. The gradient of (3), $\partial L / \partial \log k$, must flow through the pipeline into the operator.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Shapes: n = 2,000 labelled scenes, B = 5, log k (5,)
class BandInversion(eqx.Module):                                       # user
    log_k: jax.Array                    # z_b = −log y_b / exp(log k_b)
pipe = Sequential([JaxModelOp(BandInversion(log_k0)),
                   Lambda(weighted_mean)])                              # today
jax.tree.leaves(pipe)            # → [Sequential(...)]: one opaque leaf today
# ∂L/∂ pipe
eqx.filter_grad(loss)(pipe, y, a)          # → None: nothing to diff   today
# rebuild the pipeline inside the loss; differentiate the module instead
def loss_m(m, y, a):                                                    # user
    return loss(Sequential([JaxModelOp(m), Lambda(weighted_mean)]), y, a)
# log k ← Adam(log k, ∂L/∂ log k), 400 steps, learning rate 10⁻²
eqx.filter_value_and_grad(loss_m)(m, y, a)                              # today
# test RMSE: 3.45 at the start → 0.260 (true k: 0.262); k still ≤ 8.7 % off

The plain pipeline is opaque to JAX. Sequential and JaxModelOp are ordinary Python objects, so jax.tree.leaves sees the whole pipeline as one leaf, and Equinox’s filtered gradient finds no array to differentiate. The workaround is the one pipekit-train already uses: differentiate the module, and rebuild the pipeline inside the loss. Then 400 Adam steps bring the test RMSE from 3.45 to 0.260, as good as the true coefficients (0.262).

The coefficients themselves stay up to 8.7 % off. The labels constrain the retrieval, not each band’s absorption, and several sets of five $k_b$ give the same $\hat\alpha$. A loss on $\hat\alpha$ is the right objective for the retrieval; recovering the spectroscopy needs a loss on $y$.

Stage 3 · Control, observation and caching operators under jit (pipekit)¶

TL;DR. Put each of pipekit’s side-effecting and branching operators in a jitted pipeline and record what it does.

Problem. Under jit an operator’s _apply runs once, on a tracer, and only the JAX operations it records survive into the compiled function. An operator whose job is a side effect, a Python decision or a lookup keyed on its input’s value cannot do that job. Table 2 records what each one does.

1
2
3
4
5
6
7
8
9
10
11
# Shapes: x (5,) = (1, …, 5)
cache = Cache(Lambda(lambda x: x + 1))                                  # today
twice = Sequential([cache, Lambda(lambda x: 10 * x), cache])            # today
twice(x)                  # eager: (21, 31, 41, 51, 61)                 today
jax.jit(twice)(x)         # → (2, 3, 4, 5, 6): the first call's output  today
# key(x) = sha256(repr(x)) for a tracer: one key per shape and dtype
# replacement for Tap: runs on every call, not once per trace
jax.debug.callback(fn)    # 3 calls → 3 callbacks                       today
# replacement for Branch: both arms traced, one selected at run time
CondBranch(pred, retrieval, last_band)                  # lax.cond     proposed (C1)
# → equals the eager Branch to 10⁻¹²

Four operators are honest about it: Branch and Switch raise (TracerBoolConversionError is a subclass of ConcretizationTypeError), and Fanout and a diamond Graph simply trace. Four run at trace time only, which is what the scoping report predicted; ShapeTrace is useful that way, Tap, Profile and Sink are not, and Snapshot keeps a tracer that raises when touched later.

Cache is the dangerous one. A tracer has no tobytes, so _input_key (cache.py:97) falls back to stable_repr, and every tracer of one shape and dtype has the same representation. The second cache call in twice hits the first one’s entry and returns $x + 1$ in place of $10(x + 1) + 1$, with no error. The stored tracers then poison later traces of other functions.

Stage 4 · Operators as PyTrees, prototyped (pipekit-jax)¶

TL;DR. Prototype the scoping report’s JaxOperator and JaxSequential in a few lines, and see what the multiple inheritance gives and costs.

Problem. For $\partial L / \partial \theta$ to reach parameters $\theta$ held by operators, each operator must be a PyTree whose array fields are leaves, and so must the container.

1
2
3
4
5
6
7
8
9
10
11
12
# Shapes: B = 5; two operators, each holding log k (5,)
class JaxOperator(Operator, eqx.Module): pass                  # proposed (B1)
# MRO: Invert → JaxOperator → Operator → ConfigMixin → Module
class JaxSequential(Sequential, eqx.Module):                   # proposed (B1)
    operators: tuple
seq = JaxSequential([Invert(log_k0), Combine(log_k0)])         # proposed (B1)
jax.tree.leaves(seq)            # → 2 leaves, one per log k     proposed (B1)
# ∂L/∂θ for both operators; the tied gradient is their sum
eqx.filter_grad(loss)(seq, y, a)       # = stage 2's, to 10⁻⁹   proposed (B1)
eqx.tree_at(lambda s: s.operators[0], seq, Invert(log_k))      # proposed (B1)
seq.summary(None)               # → still a pipekit Sequential  proposed (B1)
json.dumps(Invert(log_k0).state)       # → TypeError: an Array  proposed (B2)

The multiple inheritance works as the scoping report said: Equinox’s Module registers the class as a PyTree and pipekit’s Operator keeps dispatch, summary and |. A single JaxOperator is differentiable; put in a plain Sequential it is opaque again, so the container must be a PyTree too. With both, the gradient reaches every operator, and eqx.tree_at can replace one step. Because each operator holds its own $\log k$, the two copies are separate leaves, and their gradients sum to the gradient of stage 2’s tied parameter. Tying them again is what eqx.nn.Shared is for.

One cost the report did not mention. ConfigMixin derives get_config from the __init__ parameters, and Equinox’s generated __init__ takes the array fields. So the config holds a JAX array, and state is no longer JSON. The array values belong in the registry’s weight blob, as JaxModelOp already does it, and the config should record only their shapes and dtypes (B2).

Stage 5 · What pipekit-jax ships (pipekit-jax → pipekit-experiment, pipekit-cycle)¶

TL;DR. Round-trip the trained inversion through the registry, truncate a scan’s gradient, and solve an ODE forward model with three adjoints.

Problem. Three shipped pieces, each with a check. (a) A reloaded model must equal the stored one byte for byte. (b) For a plume column refilled each hour, $c_{t+1} = 0.9\, c_t + q_t$ over $T = 24$ hours, the exact sensitivity of the last column to the first hour’s source is (4),

and truncating to the last $k$ steps must zero every earlier term. (c) For first-order loss $\dot c = -\lambda c$, $\lambda = 0.25\ \text{h}^{-1}$, one hour’s solve and its gradient are both $e^{-\lambda} = 0.7788$.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Shapes: log k (5,); q (24,); c ()
# (a) the stage 2 module, through the registry and back
h = reg.store(JaxModelOp(m), weights=JaxModelOp(m).serialize_weights())  # today
JaxModelOp(BandInversion(jnp.zeros(5))).from_registry(reg, h)           # today
# → byte-identical weights; max |k_reloaded − k̂| = 0
# (b) ∂c_T/∂q_t through the carry, kept for the last k steps only
truncated_scan(decay, 0.0, q, k=3)     # → 3 nonzero gradients    today
truncated_scan(decay, 0.0, q, k=24)    # → 24; ∂c_T/∂q_0 = 0.9²³ = 0.0886  today
# forward values identical to lax.scan
# (c) ċ = −λc over one hour
ode = DiffraxForwardModel(vector_field=lambda t, c, a: -LAM * c,
                          dt=1.0, dt0=0.01)                             # today
ode.step(100.0, 1.0)                   # → 77.88, error < 10⁻¹²     today
ode.with_adjoint(pc.DirectAdjoint())   # also Backsolve; ∂c/∂c₀ = 0.7788  today
ode.with_adjoint(pc.TruncatedAdjoint(k=3)).step(1.0, 1.0)  # → ValueError  today

All three work as documented. JaxModelOp round-trips the weights exactly through the registry, which is how stage 4’s operators should store theirs. truncated_scan keeps exactly the last $k$ gradients through the carry and leaves the forward values alone. DiffraxForwardModel satisfies pipekit-cycle’s ForwardModel and SupportsAdjoint, and the checkpointed, direct and backsolve adjoints give the same gradient on this non-stiff problem. A TruncatedAdjoint spec raises, by design: truncation is a property of a rollout, which truncated_scan handles, not of one solve.

The scoping report’s names, checked¶

Table 3 checks each name the scoping report’s examples use.

Current-state audit¶

Of twelve pieces checked, six ship and work, three ship with defects or run at trace time only, one returns wrong values, and two are missing (Table 4); Figure 2 counts the page’s calls by marker. The audit read references/geostack/pipekit at ef3c17f on 14 October 2026.

Smaller findings:

• Reports 2 and 3 said a pipeline cannot be jitted. Report 2’s exclusions (“these clash with dual-mode __call__”) and Report 3’s “jax.jit does not compose with a pipekit pipeline” repeated the claim the scoping report corrected; both now link to stage 1.

• pipekit-array’s namespace dispatch already supports jax.numpy arrays (_namespace.py); a pipekit-array operator in a jitted pipeline is pure JAX when its input is a JAX array. Not run here.

• pipekit-jax depends on pipekit-cycle in its base install, only for diffrax_model.py’s default adjoint spec. Report 16’s B moves it to the [diffrax] extra.

Workstream A · Trace guards in pipekit core¶

Decision. pipekit core stays free of JAX, but an operator that cannot do its job on a tracer says so. Detection uses sys.modules: if jax is already imported, isinstance(x, jax.core.Tracer); otherwise nothing is traced and nothing is checked. Table 5 lists the moves.

1
2
3
4
5
6
# Shapes: x (5,)
jax.jit(Sequential([cache, Lambda(lambda x: 10 * x), cache]))(x)
# → TypeError: Cache cannot key a traced value; inside jax.jit,
#   JAX's compilation cache already does this                proposed (A1)
jax.jit(Sequential([snap.at("z"), retrieval]))(y)
# → TypeError: Snapshot cannot store a traced value          proposed (A2)

Workstream B · Operators as PyTrees in pipekit-jax¶

Decision. When a project differentiates or vmaps over operator parameters, pipekit-jax adds the scoping report’s mixin. Until then the documented pattern is stage 2’s: differentiate the module, rebuild the pipeline in the loss.

1. B1 · JaxOperator and JaxSequential. JaxOperator(Operator, eqx.Module) as in stage 4; JaxSequential(Sequential, eqx.Module) with operators: tuple, keeping Sequential’s summary, | and stateful path. Config fields are eqx.field(static=True), so two pipelines with the same config share one compiled function.

2. B2 · An array policy for configs. get_config of a JaxOperator records each array field as {"shape": …, "dtype": …}, forbid_in_yaml stays false for the structure, and the values travel as the registry’s weight blob through serialize_weights and with_weights, shared with JaxModelOp.

3. No JaxGraph yet. A Graph already traces as a callable (stage 3); a PyTree Graph meets the tree-not-DAG problem the scoping report described and waits for a use.

1
2
3
4
5
# Shapes: log k (5,) per operator
seq = pkj.JaxSequential([pkj.Invert(log_k0), pkj.Combine(log_k0)])  # proposed (B1)
g = eqx.filter_grad(loss)(seq, y, a)    # reaches both log k     proposed (B1)
reg.store(seq, weights=seq.serialize_weights())                  # proposed (B2)
json.dumps(seq.state)                   # shapes and dtypes only proposed (B2)

Workstream C · JAX-native control and observation¶

Decision. Branching on a traced value and observing at run time are operators in pipekit-jax, with pipekit’s names and JAX’s semantics.

1. C1 · CondBranch(predicate, if_true, if_false) and CondSwitch(index, branches) on jax.lax.cond and jax.lax.switch. Both arms are traced and must return the same shape and dtype; the constructor cannot check that, so the first call does.

2. C2 · DebugTap(fn) on jax.debug.callback, which runs fn on concrete values on every call (three calls, three callbacks in stage 3). Its ordering across devices is not guaranteed; the docstring says so.

Workstream D · Transformation wrappers¶

Decision. Not now. jax.vmap, jax.checkpoint and jax.lax.scan apply to a pipeline directly, as stage 1 does with vmap. A VMapOver(op), Checkpoint(op) or ScanSequential(op, n) operator adds a name and a config to the transformation, which matters only when a pipeline that uses one must round-trip through YAML. The gate is a project that needs that.

Workstream E · Docs¶

1. E1 · Build time and trace time. One page in pipekit’s docs: what runs once at trace time (ShapeTrace, Tap, a Lambda’s Python), what survives into the compiled function, why Cache refuses a tracer, and stage 2’s pattern for differentiating a module inside a pipeline.

2. E2 · pipekit-jax’s README and install notes. Drop the link to docs/design/ (or add the notes), remove orbax-checkpoint from the install notes, and drop jaxtyping from the dependencies or use it.

Testing gates¶

Table 6 gives each workstream’s gate.

Sequencing and risks¶

Figure 3 orders the work. Phase 1 is small and protects users now: A1 is a few lines and a test, and it removes the one silent wrong answer. Phase 2 starts when a project fits operator parameters, as stage 2 does; until then stage 2’s pattern works. Phase 3 waits for a pipeline that must round-trip a transformation through YAML.

• Equinox’s Module and ConfigMixin disagree about __init__. The mixin reads __init__'s parameters; Equinox generates __init__ from fields. Stage 4 shows they coexist, but a custom __init__ on a JaxOperator must keep the field names, or get_config breaks. B1’s test covers both.

• Static fields and retracing. A config field left as a leaf retraces on every new value, and an array marked static breaks hashing. B1’s compile-once check catches the first; Equinox raises on the second.

• The workaround may be enough. If every project trains one module per pipeline, as pipekit-train does, B never needs to ship. That is the open question below.

Open questions¶

1. Does any project need gradients through more than one operator? Stage 2’s rebuild pattern covers one learned module. Workstream B is worth it only for pipelines with several learned steps, or for vmapping over whole pipelines (the scoping report’s hyperparameter sweeps, 7.6).

2. Should Cache refuse tracers or pass them through? Refusing (A1) is explicit; passing through to the inner operator would make a pipeline with a Cache jittable unchanged, at the cost of a cache that silently does nothing under jit.

3. Where do the PyTree operators live if geotoolz needs them? A JAX version of a GeoTensor operator (the scoping report’s jax_geotoolz) would subclass JaxOperator. Whether that is a geotoolz extra or a new package is open, and nothing needs it yet.

4. JaxModelOp and JaxOperator: one class or two? A wrapper keeps any module unchanged; a mixin makes the operator itself the module. Report 16 settles on JaxModelOp as the one wrapper; B1 would add the mixin beside it, not replace it.

Upcoming work¶

Table 7 lists the follow-ups.

References¶

Sources¶

• The scoping report, “Report 5 — pipekit-jax: a future-direction analysis” Johnson, 2026.

• pipekit 0.0.2 and pipekit-jax 0.0.2 in references/geostack/pipekit at ef3c17f: source, tests, README, CHANGELOG and docs.

• JAX Bradbury et al., 2018, Equinox Kidger & Garcia, 2021 and diffrax Kidger, 2021, at the versions in the header.

• The running example’s notebook, described in Section ??.