Skip to content

pipekit-jax

JAX / Equinox carrier integration for pipekit. Ships JaxModelOp (the public counterpart of the in-package pipekit_train.adapters.equinox.EquinoxModelOp, adding byte-identical weight round-trip through pipekit_experiment.ModelRegistry via the registry's weights blob), DiffraxForwardModel (a pipekit_cycle.ForwardModel behind the [diffrax] extra), and the adjoint interpreters to_diffrax_adjoint / truncated_scan.

Pytree containers

Importing pipekit_jax registers pipekit.Sequential, pipekit.Graph and JaxModelOp as JAX pytrees, so jax.grad / eqx.filter_grad / jax.vmap see through a pipeline to the leaves of the operators it holds.

pipekit_jax.pytree

Register pipekit's composition containers as JAX pytrees.

Importing pipekit_jax registers pipekit.Sequential, pipekit.Graph and pipekit_jax.JaxModelOp with jax.tree_util, so JAX transformations see through a pipeline to the leaves of the operators it holds:

  • Sequential — children are its operators, in order.
  • Graph — children are the node operators in topological order; the topology (inputs, node wiring, outputs) is static aux data that holds no operators or parameters.
  • JaxModelOp — the child is its wrapped eqx.Module.

Without the registration a container is a single opaque leaf, and a learnable coefficient on an operator inside it is a constant to jax.grad / eqx.filter_grad. pipekit core stays JAX-free: the registration is an import side effect of this package, like enabling a jax.config flag.

What the leaves are depends on the operators. An operator that is itself a pytree (an eqx.Module subclass, or JaxModelOp via its module) contributes its own leaves. A plain pipekit.Operator is one opaque leaf — fine under eqx.filter_jit / eqx.filter_grad (non-array leaves are static there), but not as a direct argument to jax.jit.

Caveats:

  • Registration is by exact type. Subclasses of Sequential / Graph are not registered automatically.
  • Operator sharing is not preserved. An operator instance used by two graph nodes flattens to two children; after an unflatten with new leaves they are two independent objects. (Equinox modules behave the same way.)
  • Other containers (Fanout, Branch, Switch, Cache, Try, the parallel maps) are not registered yet.
Example

import equinox as eqx, jax, jax.numpy as jnp import pipekit_jax # registers the containers from pipekit import Operator, Sequential class Scale(Operator, eqx.Module): ... factor: jax.Array ... def _apply(self, x): ... return self.factor * x pipe = Sequential([Scale(jnp.asarray(2.0)), Scale(jnp.asarray(3.0))]) len(jax.tree.leaves(pipe)) 2

JaxOperator

Base class for operators that are both a pipekit.Operator and an eqx.Module: JSON-safe config, pytree leaves, and the weight-blob round trip. Use check_jax_operator in a library's test suite to assert the contract (eager ≡ jit, batching, visible leaves, a config round trip that keeps learnable fields learnable).

pipekit_jax.operator

JaxOperator — a pipekit Operator that is also an eqx.Module.

Domain libraries (finitevolX, spectraldiffx, somax, …) need operators that are both:

  • pipekit operators — dual-mode __call__ (eager / graph node), | composition, and a JSON state / from_state round trip;
  • Equinox modules — array fields are pytree leaves, so the operator works under jax.jit / jax.grad / jax.vmap, and config-only fields are eqx.field(static=True).

The two compose because Equinox generates a dataclass __init__ from the fields, and that is exactly the signature pipekit's ConfigMixin reads to build get_config(). JaxOperator makes the combination explicit, coerces array fields to JSON in get_config(), and adds the WeightsMixin byte round trip.

Writing one::

class Diffuse(JaxOperator):
    '''u ↦ u + κ ∇²u on a periodic 1-D grid (κ learnable).'''

    dx: float = eqx.field(static=True)        # config, not a leaf
    kappa: jax.Array = eqx.field(converter=jnp.asarray)   # leaf

    def _apply(self, u):
        lap = (jnp.roll(u, -1) - 2 * u + jnp.roll(u, 1)) / self.dx**2
        return u + self.kappa * lap

Learnable fields: give them converter=jnp.asarray. A JSON reload passes a Python float back into the constructor, and a Python float is not an inexact array — without the converter the reloaded parameter is silently treated as a constant by eqx.filter_grad. pipekit_jax.testing.check_jax_operator catches this.

Compiling: use eqx.filter_jit(op) (or jax.jit on a function that takes op as an argument). jax.jit(op) hashes the callable, and an eqx.Module with array fields is unhashable.

JaxOperator

Bases: Operator, WeightsMixin, Module

Base class for pipekit operators whose fields are Equinox fields.

Subclasses declare fields as on any eqx.Module and implement _apply. Operator comes first in the MRO, so pipekit's dual-mode __call__ is the one that runs (Equinox defines no __call__).

  • Config: get_config() is auto-derived from the generated __init__ and coerced to JSON builtins (jsonable), so op.state survives json.dumps and Operator.from_state. Real numbers, strings, booleans, (nested) lists / tuples and str-keyed dicts of them round-trip, and so do nested operators (their state records, which from_state rebuilds); op.state raises TypeError for complex values, and for array fields whose converter would not restore their dtype (JSON drops dtypes: a float16 field needs a casting converter). Operators holding non-array Python objects (callables, modules without a JSON form) should set forbid_in_yaml = True and persist through serialize_weights / with_weights.
  • Leaves: inexact-array fields are differentiable; mark config-only fields eqx.field(static=True); give learnable fields converter=jnp.asarray (see the module docstring).
  • Weights: serialize_weights / with_weights / from_registry persist the operator's own array leaves.
Source code in packages/pipekit-jax/src/pipekit_jax/operator.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
class JaxOperator(Operator, WeightsMixin, eqx.Module):
    """Base class for pipekit operators whose fields are Equinox fields.

    Subclasses declare fields as on any ``eqx.Module`` and implement
    ``_apply``. ``Operator`` comes first in the MRO, so pipekit's
    dual-mode ``__call__`` is the one that runs (Equinox defines no
    ``__call__``).

    - **Config**: ``get_config()`` is auto-derived from the generated
      ``__init__`` and coerced to JSON builtins (`jsonable`), so
      ``op.state`` survives ``json.dumps`` and ``Operator.from_state``.
      Real numbers, strings, booleans, (nested) lists / tuples and
      str-keyed dicts of them round-trip, and so do nested operators
      (their ``state`` records, which ``from_state`` rebuilds);
      ``op.state`` raises ``TypeError`` for complex values, and for
      array fields whose converter would not restore their dtype (JSON
      drops dtypes: a ``float16`` field needs a casting converter).
      Operators holding non-array Python objects (callables, modules
      without a JSON form) should set ``forbid_in_yaml = True`` and
      persist through `serialize_weights` / `with_weights`.
    - **Leaves**: inexact-array fields are differentiable; mark
      config-only fields ``eqx.field(static=True)``; give learnable
      fields ``converter=jnp.asarray`` (see the module docstring).
    - **Weights**: `serialize_weights` / `with_weights` /
      `from_registry` persist the operator's own array leaves.
    """

    def get_config(self) -> dict[str, Any]:
        """JSON-builtin constructor config (arrays → scalars / lists)."""
        return {k: jsonable(v) for k, v in super().get_config().items()}

    @property
    def state(self) -> dict[str, Any]:
        """JSON-serialisable state record, validated for a round trip.

        Raises:
            TypeError: if a config value cannot survive ``json.dumps`` →
                ``Operator.from_state`` (complex numbers, dicts) and the
                class is not marked ``forbid_in_yaml``.
        """
        record = super().state
        if type(self).forbid_in_yaml:
            # Not round-trippable by design, but still JSON-encodable so a
            # registry can store the record next to the weight blob.
            record["config"] = {
                k: v if _json_safe(v) else _debug_repr(v)
                for k, v in record["config"].items()
            }
            return record
        if not type(self).forbid_in_yaml:
            bad = [k for k, v in record["config"].items() if not _round_trips(v)]
            if bad:
                raise TypeError(
                    f"{type(self).__name__}: config field(s) {bad} cannot round-trip "
                    "through JSON and Operator.from_state (complex values, and "
                    "dicts keyed exactly {'class', 'config'}, are not supported). "
                    "Store them as real numbers or lists of pairs, or set "
                    "`forbid_in_yaml = True` and persist arrays with "
                    "serialize_weights()."
                )
            lossy = self._dtype_lossy_fields()
            if lossy:
                raise TypeError(
                    f"{type(self).__name__}: array field(s) {lossy} would reload "
                    "with a different dtype: JSON drops dtypes and the field's "
                    "converter does not restore it. Give the field a casting "
                    "converter, e.g. eqx.field(converter=lambda x: "
                    "jnp.asarray(x, jnp.float16)), or set `forbid_in_yaml = True` "
                    "and persist arrays with serialize_weights()."
                )
        return record

    def _dtype_lossy_fields(self) -> list[str]:
        """Array config fields whose JSON round trip changes their dtype.

        Rebuilds each array field the way ``from_state`` does (its JSON
        value passed through the field's converter) and compares dtypes.
        A field without a converter reloads as a Python scalar / list;
        `check_jax_operator` reports that case as a lost learnable leaf.
        """
        lossy = []
        for f in dataclasses.fields(self):
            converter = f.metadata.get("converter")
            if not f.init or converter is None:
                continue
            value = getattr(self, f.name)
            if not eqx.is_array(value):
                continue
            rebuilt = converter(jsonable(value))
            if not eqx.is_array(rebuilt) or rebuilt.dtype != value.dtype:
                lossy.append(f.name)
        return lossy

state property

JSON-serialisable state record, validated for a round trip.

Raises:

Type Description
TypeError

if a config value cannot survive json.dumps → Operator.from_state (complex numbers, dicts) and the class is not marked forbid_in_yaml.

get_config()

JSON-builtin constructor config (arrays → scalars / lists).

Source code in packages/pipekit-jax/src/pipekit_jax/operator.py 
155
156
157
def get_config(self) -> dict[str, Any]:
    """JSON-builtin constructor config (arrays → scalars / lists)."""
    return {k: jsonable(v) for k, v in super().get_config().items()}

jsonable(value)

Coerce a config value to JSON builtins.

Arrays (JAX or NumPy) become a Python scalar (0-d) or nested lists; tuples (including named tuples) become plain tuples and lists plain lists, converted recursively — JSON drops the subclass anyway; dicts are converted recursively; everything else is returned unchanged.

Source code in packages/pipekit-jax/src/pipekit_jax/operator.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def jsonable(value: Any) -> Any:
    """Coerce a config value to JSON builtins.

    Arrays (JAX or NumPy) become a Python scalar (0-d) or nested lists;
    tuples (including named tuples) become plain tuples and lists plain
    lists, converted recursively — JSON drops the subclass anyway; dicts
    are converted recursively; everything else is returned unchanged.
    """
    if isinstance(value, np.ndarray | np.generic) or eqx.is_array(value):
        arr = np.asarray(value)
        return arr.item() if arr.ndim == 0 else arr.tolist()
    if isinstance(value, tuple):
        return tuple(jsonable(v) for v in value)
    if isinstance(value, list):
        return [jsonable(v) for v in value]
    if isinstance(value, dict):
        return {k: jsonable(v) for k, v in value.items()}
    return value

pipekit_jax.testing

Contract checks for JAX operators built on JaxOperator.

Call check_jax_operator from a library's test suite on every operator it ships::

def test_contract():
    check_jax_operator(Diffuse(dx=0.1, kappa=0.5), jnp.sin(x))

check_jax_operator(op, *inputs, vmap=True)

Assert that op behaves as a JAX-native pipekit operator.

Checks, in order:

  1. op is a pipekit Operator and an eqx.Module.
  2. Eager ≡ jit: op(*inputs) equals eqx.filter_jit of the same call, with op passed as an argument (so its leaves are traced).
  3. Batching (when vmap): stacking every input array along a new leading axis and jax.vmap-ing gives the stacked eager outputs. Non-array inputs (strings, callables, metadata) are held static.
  4. Leaves visible: every inexact-array field of op is a leaf of jax.tree.leaves(op), and eqx.filter_grad of a scalar loss returns a gradient for each of them.
  5. Config round trip (unless forbid_in_yaml): op.state survives json.dumps → Operator.from_state; the rebuilt operator has the same number of inexact-array leaves (a learnable field that came back as a Python float fails here), the same array leaves (shape / dtype), and the same outputs. Static fields may change container type (a tuple reloads as a list). For an operator whose config nests operators ({"class", "config"} payloads), the children are taken from the live operator and the rest is rebuilt from JSON, so local fields are still checked.
  6. No static field holds an inexact array (it could never train).

Parameters:

Name Type Description Default
op Any

The operator under test.

 required
*inputs Any

Example positional inputs for op.

 ()
vmap bool

Also check batching over a new leading axis.

 True

Raises:

Type Description
AssertionError

on the first violated property.
Source code in packages/pipekit-jax/src/pipekit_jax/testing.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def check_jax_operator(op: Any, *inputs: Any, vmap: bool = True) -> None:
    """Assert that ``op`` behaves as a JAX-native pipekit operator.

    Checks, in order:

    1. ``op`` is a pipekit ``Operator`` and an ``eqx.Module``.
    2. Eager ≡ jit: ``op(*inputs)`` equals ``eqx.filter_jit`` of the same
       call, with ``op`` passed as an argument (so its leaves are traced).
    3. Batching (when ``vmap``): stacking every input array along a new
       leading axis and ``jax.vmap``-ing gives the stacked eager outputs.
       Non-array inputs (strings, callables, metadata) are held static.
    4. Leaves visible: every inexact-array field of ``op`` is a leaf of
       ``jax.tree.leaves(op)``, and ``eqx.filter_grad`` of a scalar loss
       returns a gradient for each of them.
    5. Config round trip (unless ``forbid_in_yaml``): ``op.state``
       survives ``json.dumps`` → ``Operator.from_state``; the rebuilt
       operator has the same number of inexact-array leaves (a learnable
       field that came back as a Python float fails here), the same
       array leaves (shape / dtype), and the same outputs. Static fields
       may change container type (a tuple reloads as a list). For an
       operator whose config nests operators (``{"class", "config"}``
       payloads), the children are taken from the live operator and the
       rest is rebuilt from JSON, so local fields are still checked.
    6. No ``static`` field holds an inexact array (it could never train).

    Args:
        op: The operator under test.
        *inputs: Example positional inputs for ``op``.
        vmap: Also check batching over a new leading axis.

    Raises:
        AssertionError: on the first violated property.
    """
    assert isinstance(op, Operator), f"{type(op).__name__} is not a pipekit Operator"
    assert isinstance(op, eqx.Module), f"{type(op).__name__} is not an eqx.Module"

    expected = op(*inputs)
    jitted = eqx.filter_jit(lambda o, *xs: o(*xs))(op, *inputs)
    _assert_tree_allclose(jitted, expected, "eager vs filter_jit")

    if vmap:
        # Batch only the array leaves; static inputs (strings, callables,
        # metadata) are closed over unbatched.
        dynamic, static = eqx.partition(inputs, eqx.is_array)
        if jax.tree.leaves(dynamic):
            batched_dynamic = jax.tree.map(lambda x: jnp.stack([x, x]), dynamic)
            batched = jax.vmap(lambda d: op(*eqx.combine(d, static)))(batched_dynamic)
            _assert_tree_allclose(
                batched,
                jax.tree.map(lambda y: jnp.stack([y, y]), expected),
                "vmap over a leading axis",
            )

    hidden = [
        f.name
        for f in dataclasses.fields(op)
        if f.metadata.get("static")
        and any(eqx.is_inexact_array(x) for x in jax.tree.leaves(getattr(op, f.name)))
    ]
    assert not hidden, (
        f"field(s) {hidden} hold inexact arrays but are static, so they are not "
        "pytree leaves and can never be trained — drop static=True"
    )
    params = [leaf for leaf in jax.tree.leaves(op) if eqx.is_inexact_array(leaf)]
    if params:

        def loss(o: Any) -> Any:
            out = o(*inputs)
            return sum(jnp.sum(jnp.real(y)) for y in jax.tree.leaves(out))

        grads = eqx.filter_grad(loss)(op)
        grad_leaves = [g for g in jax.tree.leaves(grads) if eqx.is_array(g)]
        assert len(grad_leaves) == len(params), (
            f"filter_grad returned {len(grad_leaves)} gradients for "
            f"{len(params)} inexact-array leaves"
        )

    if type(op).forbid_in_yaml:
        return
    state = json.loads(json.dumps(op.state))
    init_fields = {f.name for f in dataclasses.fields(op) if f.init}
    if not set(state["config"]) <= init_fields:
        # A custom get_config whose keys are not constructor fields (one
        # that groups children under a name, say) is checked for its
        # encoding only.
        return
    # Nested children rebuild from their own state records, so the whole
    # operator, local fields and children alike, comes back from JSON.
    rebuilt = Operator.from_state(state)
    rebuilt_params = [
        leaf for leaf in jax.tree.leaves(rebuilt) if eqx.is_inexact_array(leaf)
    ]
    assert len(rebuilt_params) == len(params), (
        f"config round trip turned {len(params) - len(rebuilt_params)} learnable "
        "field(s) into constants — declare them with "
        "eqx.field(converter=jnp.asarray)"
    )
    # Compare the array leaves, not the treedef: static fields are part of
    # the treedef, and JSON legitimately turns a static tuple into a list.
    assert _array_signature(rebuilt) == _array_signature(op), (
        "config round trip changed the operator's array leaves (shape / dtype)"
    )
    _assert_tree_allclose(rebuilt(*inputs), expected, "config round trip")

pipekit_jax.weights

Shared weight-blob round trip for pipekit-jax operators.

WeightsMixin gives an operator serialize_weights / with_weights / from_registry: the array leaves of a pytree are written with equinox.tree_serialise_leaves and read back into a skeleton — an instance with the same pytree structure, whose non-array leaves (activations, dtypes, sizes, static config) cannot be recovered from bytes alone.

Equinox loads leaves positionally, so a skeleton whose fields are in a different order, or a different class with the same field names — but the same number of leaves and compatible shapes — would silently receive weights it was never trained with. The blob therefore starts with a signature header (tree_signature): the node types of the tree (module.QualName of every pytree node, from the root down, with the structural aux data of non-Equinox nodes such as a pipekit.Graph's wiring) and the key path, shape and dtype of every array leaf. Only array leaves are stored; every non-array value comes from the skeleton. with_weights compares the signature with the skeleton before loading. Node types are read from the treedef, not its repr, and aux data is encoded by _stable, so nothing process-specific (memory addresses) enters the signature.

Blob layout::

b"PKJXW1\n" | header length (8 bytes, little endian) | JSON signature
           | equinox.tree_serialise_leaves payload

Blobs without the magic prefix (written before the header existed) are still accepted and loaded positionally.

Two hooks pick which pytree is persisted:

  • _weights_tree() — the tree to serialise (default: self).
  • _from_weights_tree(tree) — wrap a deserialised tree back into an operator (default: return tree).

JaxModelOp persists its wrapped module; JaxOperator persists itself.

WeightsMixin

serialize_weights / with_weights / from_registry.

Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
class WeightsMixin:
    """``serialize_weights`` / ``with_weights`` / ``from_registry``."""

    def _weights_tree(self) -> Any:
        return self

    def _from_weights_tree(self, tree: Any) -> Any:
        return tree

    def serialize_weights(self) -> bytes:
        """Serialise the persisted tree's array leaves to bytes.

        A signature header (key path, shape and dtype of every array
        leaf) precedes the ``equinox.tree_serialise_leaves`` payload, so
        `with_weights` can refuse a skeleton with a different layout.
        """
        header = json.dumps(tree_signature(self._weights_tree())).encode()
        buf = io.BytesIO()
        buf.write(_MAGIC)
        buf.write(len(header).to_bytes(8, "little"))
        buf.write(header)
        # Array leaves only: non-array values (a Python-float temperature,
        # a flag) are configuration and come from the skeleton.
        eqx.tree_serialise_leaves(buf, eqx.filter(self._weights_tree(), eqx.is_array))
        return buf.getvalue()

    def with_weights(self, blob: bytes) -> Any:
        """Return a new operator with weights loaded from ``blob``.

        ``self`` is the **skeleton** / template — its non-array leaves
        (activations, dtypes, layer sizes) determine the structure into
        which the persisted bytes are loaded. The bytes must come from a
        `serialize_weights` call on an operator with the same pytree
        structure.

        Args:
            blob: Bytes produced by `serialize_weights`.

        Returns:
            A fresh operator carrying the persisted weight values.

        Raises:
            TypeError: if ``blob`` is not bytes.
            ValueError: if the blob's signature (node types, key paths,
                shapes, dtypes) differs from the skeleton's, or — for a legacy
                blob without a signature — if it holds more array leaves
                than the skeleton requests.
            RuntimeError: from newer equinox versions when a
                deserialised leaf's shape/dtype doesn't match the
                skeleton (``eqx.tree_deserialise_leaves``).
        """
        name = type(self).__name__
        if not isinstance(blob, (bytes, bytearray)):
            raise TypeError(
                f"{name}.with_weights(blob): blob must be bytes "
                f"(produced by serialize_weights); got {type(blob).__name__}."
            )
        skeleton = self._weights_tree()
        buf = io.BytesIO(blob)
        if bytes(blob[: len(_MAGIC)]) == _MAGIC:
            buf.seek(len(_MAGIC))
            size = int.from_bytes(buf.read(8), "little")
            saved = json.loads(buf.read(size))
            expected = tree_signature(skeleton)
            if saved["nodes"] != expected["nodes"]:
                raise ValueError(
                    f"{name}.with_weights: the blob's tree structure does not "
                    "match the skeleton: the pytree node types differ (saved "
                    f"root {saved['nodes'][0]!r}, skeleton root "
                    f"{expected['nodes'][0]!r}). Load weights only into the "
                    "class and structure they were saved from."
                )
            saved_leaves, expected_leaves = saved["leaves"], expected["leaves"]
            if saved_leaves != expected_leaves:
                diff = next(
                    (
                        (a, b)
                        for a, b in zip(saved_leaves, expected_leaves, strict=False)
                        if a != b
                    ),
                    (
                        f"{len(saved_leaves)} array leaves",
                        f"{len(expected_leaves)} array leaves",
                    ),
                )
                raise ValueError(
                    f"{name}.with_weights: the blob's tree structure does not "
                    f"match the skeleton (saved {diff[0]!r}, skeleton "
                    f"{diff[1]!r}). Build the skeleton with the same structure, "
                    "field order and shapes as the operator that was serialised."
                )
            arrays = eqx.tree_deserialise_leaves(
                buf, eqx.filter(skeleton, eqx.is_array)
            )
            tree = eqx.combine(arrays, skeleton)
        else:
            # Legacy blob: written with Equinox's default filter.
            tree = eqx.tree_deserialise_leaves(buf, skeleton)
        unread = len(blob) - buf.tell()
        if unread:
            # The skeleton asked for fewer leaves than were saved (e.g. an
            # optional bias that is None in the skeleton): a partial restore
            # would silently drop weights.
            raise ValueError(
                f"{name}.with_weights: {unread} bytes of the blob "
                "were not consumed — the saved tree has more array leaves than "
                "the skeleton. Build the skeleton with the same structure as the "
                "operator that was serialised."
            )
        return self._from_weights_tree(tree)

    def from_registry(self, registry: Any, ref: str) -> Any:
        """Reload weights stored under ``ref`` into a copy of ``self``.

        The inverse of::

            registry.store(trained_op, weights=trained_op.serialize_weights())

        Args:
            registry: Anything satisfying
                ``pipekit_experiment.ModelRegistry`` (which includes
                ``load_weights(ref) -> bytes | None``).
            ref: Hash or tag identifying the stored model.

        Returns:
            A new operator with the persisted weights loaded into a
            copy of ``self``.

        Raises:
            TypeError: if ``registry`` doesn't expose a way to read
                the raw weight bytes.
            ValueError: if no weights are stored, or they don't match
                ``self``'s structure.
        """
        name = type(self).__name__
        loader = getattr(registry, "load_weights", None)
        if loader is None:
            raise TypeError(
                f"{name}.from_registry: registry must satisfy the "
                "ModelRegistry protocol (missing `load_weights`). Got "
                f"{type(registry).__name__}."
            )
        blob = loader(ref)
        if blob is None:
            raise ValueError(
                f"{name}.from_registry: no weights blob stored "
                f"for ref={ref!r}. The model was registered without "
                "weights bytes — pass weights=op.serialize_weights() "
                "to registry.store() when persisting it."
            )
        return self.with_weights(blob)

from_registry(registry, ref)

Reload weights stored under ref into a copy of self.

The inverse of::

registry.store(trained_op, weights=trained_op.serialize_weights())

Parameters:

Name Type Description Default
registry Any

Anything satisfying pipekit_experiment.ModelRegistry (which includes load_weights(ref) -> bytes | None).

 required
ref str

Hash or tag identifying the stored model.

 required

Returns:

Type Description
Any

A new operator with the persisted weights loaded into a
Any

copy of self.

Raises:

Type Description
TypeError

if registry doesn't expose a way to read the raw weight bytes.
ValueError

if no weights are stored, or they don't match self's structure.
Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def from_registry(self, registry: Any, ref: str) -> Any:
    """Reload weights stored under ``ref`` into a copy of ``self``.

    The inverse of::

        registry.store(trained_op, weights=trained_op.serialize_weights())

    Args:
        registry: Anything satisfying
            ``pipekit_experiment.ModelRegistry`` (which includes
            ``load_weights(ref) -> bytes | None``).
        ref: Hash or tag identifying the stored model.

    Returns:
        A new operator with the persisted weights loaded into a
        copy of ``self``.

    Raises:
        TypeError: if ``registry`` doesn't expose a way to read
            the raw weight bytes.
        ValueError: if no weights are stored, or they don't match
            ``self``'s structure.
    """
    name = type(self).__name__
    loader = getattr(registry, "load_weights", None)
    if loader is None:
        raise TypeError(
            f"{name}.from_registry: registry must satisfy the "
            "ModelRegistry protocol (missing `load_weights`). Got "
            f"{type(registry).__name__}."
        )
    blob = loader(ref)
    if blob is None:
        raise ValueError(
            f"{name}.from_registry: no weights blob stored "
            f"for ref={ref!r}. The model was registered without "
            "weights bytes — pass weights=op.serialize_weights() "
            "to registry.store() when persisting it."
        )
    return self.with_weights(blob)

serialize_weights()

Serialise the persisted tree's array leaves to bytes.

A signature header (key path, shape and dtype of every array leaf) precedes the equinox.tree_serialise_leaves payload, so with_weights can refuse a skeleton with a different layout.

Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def serialize_weights(self) -> bytes:
    """Serialise the persisted tree's array leaves to bytes.

    A signature header (key path, shape and dtype of every array
    leaf) precedes the ``equinox.tree_serialise_leaves`` payload, so
    `with_weights` can refuse a skeleton with a different layout.
    """
    header = json.dumps(tree_signature(self._weights_tree())).encode()
    buf = io.BytesIO()
    buf.write(_MAGIC)
    buf.write(len(header).to_bytes(8, "little"))
    buf.write(header)
    # Array leaves only: non-array values (a Python-float temperature,
    # a flag) are configuration and come from the skeleton.
    eqx.tree_serialise_leaves(buf, eqx.filter(self._weights_tree(), eqx.is_array))
    return buf.getvalue()

with_weights(blob)

Return a new operator with weights loaded from blob.

self is the skeleton / template — its non-array leaves (activations, dtypes, layer sizes) determine the structure into which the persisted bytes are loaded. The bytes must come from a serialize_weights call on an operator with the same pytree structure.

Parameters:

Name Type Description Default
blob bytes

Bytes produced by serialize_weights.

 required

Returns:

Type Description
Any

A fresh operator carrying the persisted weight values.

Raises:

Type Description
TypeError

if blob is not bytes.
ValueError

if the blob's signature (node types, key paths, shapes, dtypes) differs from the skeleton's, or — for a legacy blob without a signature — if it holds more array leaves than the skeleton requests.
RuntimeError

from newer equinox versions when a deserialised leaf's shape/dtype doesn't match the skeleton (eqx.tree_deserialise_leaves).
Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def with_weights(self, blob: bytes) -> Any:
    """Return a new operator with weights loaded from ``blob``.

    ``self`` is the **skeleton** / template — its non-array leaves
    (activations, dtypes, layer sizes) determine the structure into
    which the persisted bytes are loaded. The bytes must come from a
    `serialize_weights` call on an operator with the same pytree
    structure.

    Args:
        blob: Bytes produced by `serialize_weights`.

    Returns:
        A fresh operator carrying the persisted weight values.

    Raises:
        TypeError: if ``blob`` is not bytes.
        ValueError: if the blob's signature (node types, key paths,
            shapes, dtypes) differs from the skeleton's, or — for a legacy
            blob without a signature — if it holds more array leaves
            than the skeleton requests.
        RuntimeError: from newer equinox versions when a
            deserialised leaf's shape/dtype doesn't match the
            skeleton (``eqx.tree_deserialise_leaves``).
    """
    name = type(self).__name__
    if not isinstance(blob, (bytes, bytearray)):
        raise TypeError(
            f"{name}.with_weights(blob): blob must be bytes "
            f"(produced by serialize_weights); got {type(blob).__name__}."
        )
    skeleton = self._weights_tree()
    buf = io.BytesIO(blob)
    if bytes(blob[: len(_MAGIC)]) == _MAGIC:
        buf.seek(len(_MAGIC))
        size = int.from_bytes(buf.read(8), "little")
        saved = json.loads(buf.read(size))
        expected = tree_signature(skeleton)
        if saved["nodes"] != expected["nodes"]:
            raise ValueError(
                f"{name}.with_weights: the blob's tree structure does not "
                "match the skeleton: the pytree node types differ (saved "
                f"root {saved['nodes'][0]!r}, skeleton root "
                f"{expected['nodes'][0]!r}). Load weights only into the "
                "class and structure they were saved from."
            )
        saved_leaves, expected_leaves = saved["leaves"], expected["leaves"]
        if saved_leaves != expected_leaves:
            diff = next(
                (
                    (a, b)
                    for a, b in zip(saved_leaves, expected_leaves, strict=False)
                    if a != b
                ),
                (
                    f"{len(saved_leaves)} array leaves",
                    f"{len(expected_leaves)} array leaves",
                ),
            )
            raise ValueError(
                f"{name}.with_weights: the blob's tree structure does not "
                f"match the skeleton (saved {diff[0]!r}, skeleton "
                f"{diff[1]!r}). Build the skeleton with the same structure, "
                "field order and shapes as the operator that was serialised."
            )
        arrays = eqx.tree_deserialise_leaves(
            buf, eqx.filter(skeleton, eqx.is_array)
        )
        tree = eqx.combine(arrays, skeleton)
    else:
        # Legacy blob: written with Equinox's default filter.
        tree = eqx.tree_deserialise_leaves(buf, skeleton)
    unread = len(blob) - buf.tell()
    if unread:
        # The skeleton asked for fewer leaves than were saved (e.g. an
        # optional bias that is None in the skeleton): a partial restore
        # would silently drop weights.
        raise ValueError(
            f"{name}.with_weights: {unread} bytes of the blob "
            "were not consumed — the saved tree has more array leaves than "
            "the skeleton. Build the skeleton with the same structure as the "
            "operator that was serialised."
        )
    return self._from_weights_tree(tree)

leaf_signature(tree)

[key path, shape, dtype] of every array leaf of tree, in order.

Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
102
103
104
105
106
107
108
109
def leaf_signature(tree: Any) -> list[list[Any]]:
    """``[key path, shape, dtype]`` of every array leaf of ``tree``, in order."""
    flat, _ = jax.tree_util.tree_flatten_with_path(tree)
    return [
        [jax.tree_util.keystr(path), list(leaf.shape), str(leaf.dtype)]
        for path, leaf in flat
        if eqx.is_array(leaf)
    ]

tree_signature(tree)

Node types plus leaf_signature of tree: what a weight blob records.

Source code in packages/pipekit-jax/src/pipekit_jax/weights.py 
93
94
95
96
97
98
99
def tree_signature(tree: Any) -> dict[str, Any]:
    """Node types plus `leaf_signature` of ``tree``: what a weight blob records."""
    arrays_only = eqx.filter(tree, eqx.is_array)
    return {
        "nodes": _node_types(jax.tree.structure(arrays_only)),
        "leaves": leaf_signature(tree),
    }

ScanCycle and Rollout

jax.lax.scan-based time stepping. ScanCycle is the jit-native counterpart of pipekit_cycle.Cycle — same (carrier, state) contract and history convention, one scan regardless of n_steps, history returned by ScanCycle.run. Rollout is the carrier-only form, returning a trajectory sampled every save_every steps. checkpoint=True rematerialises in chunks on the backward pass: a ScanCycle call stores O(√n) carries instead of O(n); ScanCycle.run and Rollout store one carry per recorded / saved point plus one block, O(n / stride + stride), so they reach O(√n) only when the stride is about √n.

Measured on a 2-vCPU CPU VM (5-point diffusion step on a 128² grid):

steps compile Cycle (unrolled) compile ScanCycle grad compile Cycle grad compile ScanCycle
10 0.19 s 0.12 s 2.9 s 0.26 s
50 — — 356 s 0.28 s
100 1.04 s 0.14 s — —
1000 19.0 s 0.13 s did not finish in 10 min —
steps grad temp memory, checkpoint=False checkpoint=True
100 6.6 MiB 1.7 MiB
1000 62.9 MiB 4.4 MiB
4000 250.4 MiB 8.3 MiB

pipekit_jax.cycle

ScanCycle and Rollout — jax.lax.scan-based time stepping.

pipekit_cycle.Cycle is carrier-agnostic and steps in a Python for loop. Under jax.jit that loop unrolls: the traced program, compile time and memory all grow linearly with n_steps, and save_history appends to self.history from inside traced code. That is the right tool for host-side stepping (chunked runs with monitors between chunks) and the wrong one for anything differentiated end to end — 4DVar windows, rollout training.

  • ScanCycle — the JAX specialisation of Cycle: same (carrier, state) -> (carrier, state) contract and the same history convention, one lax.scan regardless of n_steps, and the history is returned by ScanCycle.run, never stored.
  • Rollout — the carrier-only form most PDE use needs: carrier -> trajectory sampled every save_every steps.

Both accept a pipekit operator or a pipekit_cycle.ForwardModel as the step, are JaxOperator s (so learnable leaves of the step are reachable by eqx.filter_grad), and handle arbitrary pytree carriers — dicts, eqx.Module states, custom carriers — by stacking with jax.tree.map.

EnsembleScanCycle

Bases: StatefulOperator, JaxOperator

Run a ScanCycle on every ensemble member at once, with jax.vmap.

The carrier is a pytree whose leaves share a leading member axis of size M; one compiled program advances all members.

.. code-block:: text

cⱼ ← cycle(cⱼ, s),    j = 1 … M   (in parallel)

State handling (state_axis):

  • None (default) — one state shared by all members and treated as read-only: every step of every member reads the incoming state; the state a step returns is discarded (also between the steps of the cycle), and the incoming state is returned unchanged.
  • 0 — one state per member, stacked on a leading axis like the carrier; each member threads its own state.

Not a drop-in for pipekit_cycle.EnsembleCycle when the step writes shared state: EnsembleCycle threads one state through the members sequentially (member j's output state feeds member j + 1), which a vmap cannot reproduce. Use state_axis=0 or EnsembleCycle for that.

Attributes:

Name Type Description
cycle ScanCycle

The per-member ScanCycle.
state_axis int | None

None (shared, read-only) or 0 (per member).

Raises:

Type Description
ValueError

If state_axis is not None or 0.
Source code in packages/pipekit-jax/src/pipekit_jax/cycle.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
class EnsembleScanCycle(StatefulOperator, JaxOperator):
    """Run a `ScanCycle` on every ensemble member at once, with ``jax.vmap``.

    The carrier is a pytree whose leaves share a leading member axis of
    size ``M``; one compiled program advances all members.

    .. code-block:: text

        cⱼ ← cycle(cⱼ, s),    j = 1 … M   (in parallel)

    State handling (``state_axis``):

    - ``None`` (default) — one state shared by all members and treated
      as **read-only**: every step of every member reads the incoming
      state; the state a step returns is discarded (also between the
      steps of the cycle), and the incoming state is returned unchanged.
    - ``0`` — one state per member, stacked on a leading axis like the
      carrier; each member threads its own state.

    Not a drop-in for `pipekit_cycle.EnsembleCycle` when the step
    *writes* shared state: ``EnsembleCycle`` threads one state through
    the members sequentially (member ``j``'s output state feeds member
    ``j + 1``), which a ``vmap`` cannot reproduce. Use ``state_axis=0``
    or ``EnsembleCycle`` for that.

    Attributes:
        cycle: The per-member `ScanCycle`.
        state_axis: ``None`` (shared, read-only) or ``0`` (per member).

    Raises:
        ValueError: If ``state_axis`` is not ``None`` or ``0``.
    """

    cycle: ScanCycle
    state_axis: int | None = eqx.field(static=True, default=None)

    def __check_init__(self) -> None:
        if self.state_axis not in (None, 0):
            raise ValueError(
                "EnsembleScanCycle.state_axis must be None or 0, "
                f"got {self.state_axis}."
            )

    def _read_only_cycle(self, state: Any) -> ScanCycle:
        """The cycle with its step reading ``state`` directly (not via the carry)."""
        return dataclasses.replace(
            self.cycle, step_op=_ReadOnlyState(self.cycle.step_op, state)
        )

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        # carrier: every leaf has a leading member axis [M, ...]
        if self.state_axis is None:
            cycle = self._read_only_cycle(state)
            out = jax.vmap(lambda c: cycle(c, None)[0])(carrier)
            return out, state
        return jax.vmap(self.cycle)(carrier, state)

    def run(self, carrier: Any, state: Any) -> tuple[Any, Any, Any]:
        """Like `ScanCycle.run`, per member: histories gain a leading ``M`` axis.

        Returns ``(members, state, (carriers, states))`` with
        ``carriers`` shaped ``[M, n_records, …]``.
        """
        if self.state_axis is None:
            cycle = self._read_only_cycle(state)
            c, _, hist = jax.vmap(lambda m: cycle.run(m, None))(carrier)
            return c, state, hist
        return jax.vmap(self.cycle.run)(carrier, state)

    def get_config(self) -> dict[str, Any]:
        return {"cycle": nested_config(self.cycle), "state_axis": self.state_axis}

run(carrier, state)

Like ScanCycle.run, per member: histories gain a leading M axis.

Returns (members, state, (carriers, states)) with carriers shaped [M, n_records, …].

Source code in packages/pipekit-jax/src/pipekit_jax/cycle.py 
364
365
366
367
368
369
370
371
372
373
374
def run(self, carrier: Any, state: Any) -> tuple[Any, Any, Any]:
    """Like `ScanCycle.run`, per member: histories gain a leading ``M`` axis.

    Returns ``(members, state, (carriers, states))`` with
    ``carriers`` shaped ``[M, n_records, …]``.
    """
    if self.state_axis is None:
        cycle = self._read_only_cycle(state)
        c, _, hist = jax.vmap(lambda m: cycle.run(m, None))(carrier)
        return c, state, hist
    return jax.vmap(self.cycle.run)(carrier, state)

Rollout

Bases: JaxOperator

Carrier-only rollout: c₀ -> [c₀, c_k, c_2k, …, c_n].

.. code-block:: text

cᵢ₊₁ = step(cᵢ),   i = 0 … n_steps - 1
trajectory = (c₀?, c_k, c_2k, …, c_n),   k = save_every

The trajectory is stacked along a new leading axis with jax.tree.map, so any pytree carrier works.

Attributes:

Name Type Description
step Any

A pipekit operator carrier -> carrier, or a pipekit_cycle.ForwardModel (wrapped to call step(c, dt); the wrapper is flagged forbid_in_yaml so pipekit.check_pickleable reports it).
n_steps int

Number of steps; must be a multiple of save_every.
save_every int

Record every save_every-th state.
include_initial bool

Prepend c₀ to the trajectory.
checkpoint bool

Rematerialise each save_every block on the backward pass, so grad stores carries only at save points (which the trajectory keeps anyway) — O(n / save_every + save_every) instead of O(n).

Raises:

Type Description
TypeError

If step is a StatefulOperator (use ScanCycle), or neither an operator nor a ForwardModel.
ValueError

If n_steps < 0, save_every < 1, or n_steps is not a multiple of save_every.
Source code in packages/pipekit-jax/src/pipekit_jax/cycle.py 
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
class Rollout(JaxOperator):
    """Carrier-only rollout: ``c₀ -> [c₀, c_k, c_2k, …, c_n]``.

    .. code-block:: text

        cᵢ₊₁ = step(cᵢ),   i = 0 … n_steps - 1
        trajectory = (c₀?, c_k, c_2k, …, c_n),   k = save_every

    The trajectory is stacked along a new leading axis with
    ``jax.tree.map``, so any pytree carrier works.

    Attributes:
        step: A pipekit operator ``carrier -> carrier``, or a
            `pipekit_cycle.ForwardModel` (wrapped to call ``step(c, dt)``;
            the wrapper is flagged ``forbid_in_yaml`` so
            `pipekit.check_pickleable` reports it).
        n_steps: Number of steps; must be a multiple of ``save_every``.
        save_every: Record every ``save_every``-th state.
        include_initial: Prepend ``c₀`` to the trajectory.
        checkpoint: Rematerialise each ``save_every`` block on the
            backward pass, so ``grad`` stores carries only at save points
            (which the trajectory keeps anyway) — O(n / save_every +
            save_every) instead of O(n).

    Raises:
        TypeError: If ``step`` is a `StatefulOperator` (use `ScanCycle`),
            or neither an operator nor a `ForwardModel`.
        ValueError: If ``n_steps < 0``, ``save_every < 1``, or
            ``n_steps`` is not a multiple of ``save_every``.
    """

    step: Any = eqx.field(converter=_coerce_carrier_step)
    n_steps: int = eqx.field(static=True, default=1)
    save_every: int = eqx.field(static=True, default=1)
    include_initial: bool = eqx.field(static=True, default=True)
    checkpoint: bool = eqx.field(static=True, default=False)

    def __check_init__(self) -> None:
        if self.n_steps < 0:
            raise ValueError(f"Rollout.n_steps must be >= 0, got {self.n_steps}.")
        if self.save_every < 1:
            raise ValueError(f"Rollout.save_every must be >= 1, got {self.save_every}.")
        if self.n_steps % self.save_every:
            raise ValueError(
                f"Rollout.n_steps ({self.n_steps}) must be a multiple of "
                f"save_every ({self.save_every})."
            )

    def _apply(self, carrier: Any) -> Any:
        def inner(c: Any, _: Any) -> tuple[Any, None]:
            return self.step(c), None

        def outer(c: Any, _: Any) -> tuple[Any, Any]:
            c, _ = jax.lax.scan(inner, c, None, length=self.save_every)
            return c, c

        if self.checkpoint:  # keep carries only at save points
            outer = jax.checkpoint(outer)

        n_saves = self.n_steps // self.save_every
        if n_saves:
            _, traj = jax.lax.scan(outer, carrier, None, length=n_saves)
        else:
            traj = jax.tree.map(_empty_history, carrier)
        if not self.include_initial:
            return traj
        return jax.tree.map(
            lambda c0, t: jnp.concatenate([jnp.asarray(c0)[None], t]), carrier, traj
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "step": nested_config(self.step),
            "n_steps": self.n_steps,
            "save_every": self.save_every,
            "include_initial": self.include_initial,
            "checkpoint": self.checkpoint,
        }

ScanCycle

Bases: StatefulOperator, JaxOperator

Apply step_op for n_steps inside one jax.lax.scan.

The jit-native counterpart of pipekit_cycle.Cycle.

.. code-block:: text

(cᵢ₊₁, sᵢ₊₁) = step_op(cᵢ, sᵢ),   i = 0 … n_steps - 1

Calling the operator returns (c_n, s_n) exactly like Cycle. run additionally returns the history, recorded with Cycle's convention — after step i + 1 for every i with i % history_stride == 0 — so there are ⌈n_steps / stride⌉ entries.

Attributes:

Name Type Description
step_op Any

A StatefulOperator (carrier, state) -> (carrier, state), or a pipekit_cycle.ForwardModel (wrapped so the state passes through unchanged).
n_steps int

Number of steps (>= 0; 0 is a no-op).
history_stride int

Record every stride-th step in run.
checkpoint bool

Rematerialise on the backward pass. The forward pass keeps the carry only at chunk boundaries — every ⌈√n⌉ steps when called, every history_stride steps in run — and reverse mode recomputes the steps inside a chunk. Memory for grad drops from O(n) carries to O(√n) (or O(n / stride + stride)) at the cost of one extra forward pass.
truncate int | None

If set to k, gradients flow through the carry of the trailing k steps only (the carry leaving step i is stop_gradient-ed when i < n_steps - k), matching pipekit_jax.truncated_scan: history entries recorded by run keep the gradient of their own step, as scan outputs do. Forward values are unchanged.

Raises:

Type Description
TypeError

If step_op is neither a StatefulOperator nor a ForwardModel.
ValueError

If n_steps < 0, history_stride < 1 or truncate < 1.
Example

cycle = ScanCycle(step_op=model_op, n_steps=1000, history_stride=100) carrier, state = eqx.filter_jit(cycle)(c0, s0) carrier, state, (carriers, states) = eqx.filter_jit(cycle.run)(c0, s0)

Source code in packages/pipekit-jax/src/pipekit_jax/cycle.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
class ScanCycle(StatefulOperator, JaxOperator):
    """Apply ``step_op`` for ``n_steps`` inside one ``jax.lax.scan``.

    The jit-native counterpart of `pipekit_cycle.Cycle`.

    .. code-block:: text

        (cᵢ₊₁, sᵢ₊₁) = step_op(cᵢ, sᵢ),   i = 0 … n_steps - 1

    Calling the operator returns ``(c_n, s_n)`` exactly like ``Cycle``.
    `run` additionally returns the history, recorded with ``Cycle``'s
    convention — after step ``i + 1`` for every ``i`` with
    ``i % history_stride == 0`` — so there are ``⌈n_steps / stride⌉``
    entries.

    Attributes:
        step_op: A `StatefulOperator` ``(carrier, state) -> (carrier,
            state)``, or a `pipekit_cycle.ForwardModel` (wrapped so the
            state passes through unchanged).
        n_steps: Number of steps (``>= 0``; ``0`` is a no-op).
        history_stride: Record every ``stride``-th step in `run`.
        checkpoint: Rematerialise on the backward pass. The forward pass
            keeps the carry only at chunk boundaries — every ``⌈√n⌉``
            steps when called, every ``history_stride`` steps in `run` —
            and reverse mode recomputes the steps inside a chunk. Memory
            for ``grad`` drops from O(n) carries to O(√n) (or
            O(n / stride + stride)) at the cost of one extra forward pass.
        truncate: If set to ``k``, gradients flow through the carry of
            the trailing ``k`` steps only (the carry leaving step ``i``
            is ``stop_gradient``-ed when ``i < n_steps - k``), matching
            `pipekit_jax.truncated_scan`: history entries recorded by
            `run` keep the gradient of their own step, as scan outputs
            do. Forward values are unchanged.

    Raises:
        TypeError: If ``step_op`` is neither a `StatefulOperator` nor a
            `ForwardModel`.
        ValueError: If ``n_steps < 0``, ``history_stride < 1`` or
            ``truncate < 1``.

    Example:
        >>> cycle = ScanCycle(step_op=model_op, n_steps=1000, history_stride=100)
        >>> carrier, state = eqx.filter_jit(cycle)(c0, s0)
        >>> carrier, state, (carriers, states) = eqx.filter_jit(cycle.run)(c0, s0)
    """

    step_op: Any = eqx.field(converter=_coerce_step)
    n_steps: int = eqx.field(static=True, default=1)
    history_stride: int = eqx.field(static=True, default=1)
    checkpoint: bool = eqx.field(static=True, default=False)
    truncate: int | None = eqx.field(static=True, default=None)

    def __check_init__(self) -> None:
        if self.n_steps < 0:
            raise ValueError(f"ScanCycle.n_steps must be >= 0, got {self.n_steps}.")
        if self.history_stride < 1:
            raise ValueError(
                f"ScanCycle.history_stride must be >= 1, got {self.history_stride}."
            )
        if self.truncate is not None and self.truncate < 1:
            raise ValueError(f"ScanCycle.truncate must be >= 1, got {self.truncate}.")

    def _cut(self, carry: tuple[Any, Any], i: Any) -> tuple[Any, Any]:
        """Stop gradients through the carry leaving step ``i`` (``truncate``)."""
        if self.truncate is None:
            return carry
        cut = i < self.n_steps - self.truncate
        return jax.tree.map(
            lambda x: jnp.where(cut, jax.lax.stop_gradient(x), x), carry
        )

    def _step(self, carry: tuple[Any, Any], i: Any) -> tuple[Any, Any]:
        return self._cut(self.step_op(*carry), i)

    def _apply(self, carrier: Any, state: Any) -> tuple[Any, Any]:
        chunk = max(1, math.isqrt(self.n_steps)) if self.checkpoint else None
        return _scan_steps(self._step, (carrier, state), 0, self.n_steps, chunk)

    def run(self, carrier: Any, state: Any) -> tuple[Any, Any, Any]:
        """Run the cycle and return ``(carrier, state, (carriers, states))``.

        ``carriers`` / ``states`` stack the recorded entries along a new
        leading axis of length ``⌈n_steps / history_stride⌉`` (zero when
        ``n_steps == 0``).
        """
        step, k, n = self._step, self.history_stride, self.n_steps
        carry = (carrier, state)
        if n == 0:
            return carrier, state, jax.tree.map(_empty_history, carry)

        def body(c: Any, i: Any) -> tuple[Any, None]:
            return step(c, i), None

        def group(c: Any, j: Any) -> tuple[Any, Any]:
            # one recorded step (global index j·k), then k - 1 unrecorded ones.
            # Record *before* the truncation cut, like truncated_scan's
            # outputs: an entry keeps the gradient of its own step even
            # when the carry it passes on is cut.
            recorded = self.step_op(*c)
            c = self._cut(recorded, j * k)
            if k > 1:
                c, _ = jax.lax.scan(body, c, j * k + jnp.arange(1, k))
            return c, recorded

        if self.checkpoint:  # keep carries only at group boundaries
            group = jax.checkpoint(group)

        full, rem = divmod(n, k)
        history = []
        if full:
            carry, hist = jax.lax.scan(group, carry, jnp.arange(full))
            history.append(hist)
        if rem:
            recorded = self.step_op(*carry)
            history.append(_stack([recorded]))
            carry = self._cut(recorded, full * k)
            carry = _scan_steps(step, carry, full * k + 1, rem - 1)
        hist = (
            history[0]
            if len(history) == 1
            else jax.tree.map(lambda *xs: jnp.concatenate(xs), *history)
        )
        return carry[0], carry[1], hist

    def get_config(self) -> dict[str, Any]:
        return {
            "step_op": nested_config(self.step_op),
            "n_steps": self.n_steps,
            "history_stride": self.history_stride,
            "checkpoint": self.checkpoint,
            "truncate": self.truncate,
        }

run(carrier, state)

Run the cycle and return (carrier, state, (carriers, states)).

carriers / states stack the recorded entries along a new leading axis of length ⌈n_steps / history_stride⌉ (zero when n_steps == 0).

Source code in packages/pipekit-jax/src/pipekit_jax/cycle.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def run(self, carrier: Any, state: Any) -> tuple[Any, Any, Any]:
    """Run the cycle and return ``(carrier, state, (carriers, states))``.

    ``carriers`` / ``states`` stack the recorded entries along a new
    leading axis of length ``⌈n_steps / history_stride⌉`` (zero when
    ``n_steps == 0``).
    """
    step, k, n = self._step, self.history_stride, self.n_steps
    carry = (carrier, state)
    if n == 0:
        return carrier, state, jax.tree.map(_empty_history, carry)

    def body(c: Any, i: Any) -> tuple[Any, None]:
        return step(c, i), None

    def group(c: Any, j: Any) -> tuple[Any, Any]:
        # one recorded step (global index j·k), then k - 1 unrecorded ones.
        # Record *before* the truncation cut, like truncated_scan's
        # outputs: an entry keeps the gradient of its own step even
        # when the carry it passes on is cut.
        recorded = self.step_op(*c)
        c = self._cut(recorded, j * k)
        if k > 1:
            c, _ = jax.lax.scan(body, c, j * k + jnp.arange(1, k))
        return c, recorded

    if self.checkpoint:  # keep carries only at group boundaries
        group = jax.checkpoint(group)

    full, rem = divmod(n, k)
    history = []
    if full:
        carry, hist = jax.lax.scan(group, carry, jnp.arange(full))
        history.append(hist)
    if rem:
        recorded = self.step_op(*carry)
        history.append(_stack([recorded]))
        carry = self._cut(recorded, full * k)
        carry = _scan_steps(step, carry, full * k + 1, rem - 1)
    hist = (
        history[0]
        if len(history) == 1
        else jax.tree.map(lambda *xs: jnp.concatenate(xs), *history)
    )
    return carry[0], carry[1], hist

JaxModelOp

pipekit_jax.model_op

JaxModelOp — pipekit.Operator wrapper around an eqx.Module.

Closes the v0.1 reproducibility caveat documented in pipekit-train/docs/design/boundaries.md §13.1: the in-package EquinoxModelOp stand-in shipped with v0.1 couldn't round-trip trained weights through a ModelRegistry from JSON config alone, because an eqx.Module is a JAX PyTree of arrays whose values aren't recoverable from a class name.

JaxModelOp adds two methods on top of the same wrapper surface:

  • serialize_weights() -> bytes — a leaf-signature header plus eqx.tree_serialise_leaves of the wrapped module's array leaves. Suitable for the registry's weights blob.
  • with_weights(blob: bytes) -> JaxModelOp — load the bytes back into the module's array leaves, returning a new JaxModelOp whose module is structurally identical to self.module but with the persisted weight values. self is the skeleton / template.

The skeleton-from-template requirement is honest: an eqx.Module's non-array leaves (activation functions, dtypes, hidden-size ints) are Python objects not reproducible from JSON alone. The user always needs to build a fresh module of the same shape before calling with_weights. from_registry packages the "reload weights from a stored hash into self" pattern into one call.

JaxModelOp

Bases: WeightsMixin, Operator

Wrap an eqx.Module as a pipekit.Operator with weight round-trip.

Drop-in replacement for the in-package pipekit_train.adapters.equinox.EquinoxModelOp stand-in. Same constructor signature; same _apply shape (calls the module on the input); additionally supports serialize_weights / with_weights for byte-identical persistence through any pipekit_experiment.ModelRegistry.

The wrapper is marked forbid_in_yaml = True because an eqx.Module is a PyTree of arrays — the structural config round-trips (we surface module_class in get_config) but the array values do not. The round-trip path is: registry.store(op, weights=op.serialize_weights()) to persist, then template.with_weights(blob) to reload into a fresh skeleton.

Parameters:

Name Type Description Default
module Module

An eqx.Module.

 required
Source code in packages/pipekit-jax/src/pipekit_jax/model_op.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class JaxModelOp(WeightsMixin, Operator):
    """Wrap an `eqx.Module` as a `pipekit.Operator` with weight round-trip.

    Drop-in replacement for the in-package
    ``pipekit_train.adapters.equinox.EquinoxModelOp`` stand-in.
    Same constructor signature; same ``_apply`` shape (calls the
    module on the input); additionally supports
    `serialize_weights` / `with_weights` for byte-identical
    persistence through any `pipekit_experiment.ModelRegistry`.

    The wrapper is marked ``forbid_in_yaml = True`` because an
    ``eqx.Module`` is a PyTree of arrays — the *structural* config
    round-trips (we surface ``module_class`` in `get_config`) but
    the array values do not. The round-trip path is:
    ``registry.store(op, weights=op.serialize_weights())`` to
    persist, then ``template.with_weights(blob)`` to reload into
    a fresh skeleton.

    Args:
        module: An ``eqx.Module``.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(self, module: eqx.Module) -> None:
        if not isinstance(module, eqx.Module):
            raise TypeError(
                f"JaxModelOp.module must be an eqx.Module; got {type(module).__name__}."
            )
        self.module = module

    def _apply(self, x: Any) -> Any:
        # eqx.Module is callable but ty can't see that statically.
        return cast(Any, self.module)(x)

    def get_config(self) -> dict[str, Any]:
        return {"module_class": type(self.module).__qualname__}

    # --- Weight round-trip (shared via WeightsMixin) -------------------
    # The persisted tree is the wrapped module, not the wrapper: an
    # eqx.Module's array leaves round-trip, the Operator shell doesn't.

    def _weights_tree(self) -> Any:
        return self.module

    def _from_weights_tree(self, tree: Any) -> JaxModelOp:
        return JaxModelOp(tree)

Adjoint interpreters

JAX realisations of the pipekit_cycle.adjoints specs: truncated_scan implements TruncatedAdjoint at the cycle layer (a jax.lax.scan whose carry is stop_gradient-ed outside the trailing-k window), and to_diffrax_adjoint maps specs onto diffrax.AbstractAdjoint instances for the dynamics layer.

pipekit_jax.adjoints

JAX interpreters for the pipekit adjoint-strategy specs.

Two pieces:

  • truncated_scan — a jax.lax.scan whose carry runs under stop_gradient except within the trailing k steps. The execution engine for TruncatedAdjoint at the cycle layer.
  • to_diffrax_adjoint — map a spec from pipekit_cycle.adjoints onto the corresponding diffrax.AbstractAdjoint (requires the pipekit-jax[diffrax] extra).

Specs are matched by class name and fields rather than identity, so structurally identical spec objects defined by algorithm libraries (which may not depend on pipekit) are accepted interchangeably.

to_diffrax_adjoint(spec)

Map a pipekit adjoint spec onto a diffrax.AbstractAdjoint.

Mapping (specs matched structurally by class name):

  • RecursiveCheckpointAdjoint(checkpoints) → diffrax.RecursiveCheckpointAdjoint(checkpoints) — the recommended default.
  • DirectAdjoint() → diffrax.DirectAdjoint().
  • BacksolveAdjoint() → diffrax.BacksolveAdjoint(). Warning: continuous adjoints diverge for chaotic or stiff dynamics; never use as a default.
  • ImplicitAdjoint() → diffrax.ImplicitAdjoint() — steady-state solves only.
  • TruncatedAdjoint → ValueError: truncation is a property of an outer rollout (cycle or solver iteration), not of a single ODE/SDE solve; use truncated_scan at that layer instead.

Parameters:

Name Type Description Default
spec Any

A spec instance from pipekit_cycle.adjoints or any structurally identical object. A ready-made diffrax.AbstractAdjoint passes through unchanged.

 required

Returns:

Type Description
Any

The corresponding diffrax adjoint instance.

Raises:

Type Description
ValueError

for TruncatedAdjoint or an unrecognised spec.
ModuleNotFoundError

if diffrax is not installed (install the pipekit-jax[diffrax] extra).
Source code in packages/pipekit-jax/src/pipekit_jax/adjoints.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def to_diffrax_adjoint(spec: Any) -> Any:
    """Map a pipekit adjoint spec onto a ``diffrax.AbstractAdjoint``.

    Mapping (specs matched structurally by class name):

    - ``RecursiveCheckpointAdjoint(checkpoints)`` →
      ``diffrax.RecursiveCheckpointAdjoint(checkpoints)`` — the
      recommended default.
    - ``DirectAdjoint()`` → ``diffrax.DirectAdjoint()``.
    - ``BacksolveAdjoint()`` → ``diffrax.BacksolveAdjoint()``.
      **Warning:** continuous adjoints diverge for chaotic or stiff
      dynamics; never use as a default.
    - ``ImplicitAdjoint()`` → ``diffrax.ImplicitAdjoint()`` —
      steady-state solves only.
    - ``TruncatedAdjoint`` → ``ValueError``: truncation is a property
      of an outer rollout (cycle or solver iteration), not of a single
      ODE/SDE solve; use `truncated_scan` at that layer instead.

    Args:
        spec: A spec instance from `pipekit_cycle.adjoints` or any
            structurally identical object. A ready-made
            ``diffrax.AbstractAdjoint`` passes through unchanged.

    Returns:
        The corresponding diffrax adjoint instance.

    Raises:
        ValueError: for ``TruncatedAdjoint`` or an unrecognised spec.
        ModuleNotFoundError: if diffrax is not installed (install the
            ``pipekit-jax[diffrax]`` extra).
    """
    import diffrax  # deferred: optional extra

    if isinstance(spec, diffrax.AbstractAdjoint):
        return spec
    name = type(spec).__name__
    if name == "RecursiveCheckpointAdjoint":
        return diffrax.RecursiveCheckpointAdjoint(
            checkpoints=getattr(spec, "checkpoints", None)
        )
    if name == "DirectAdjoint":
        return diffrax.DirectAdjoint()
    if name == "BacksolveAdjoint":
        return diffrax.BacksolveAdjoint()
    if name == "ImplicitAdjoint":
        return diffrax.ImplicitAdjoint()
    if name == "TruncatedAdjoint":
        raise ValueError(
            "TruncatedAdjoint applies to an outer rollout (assimilation cycle "
            "or solver iteration), not to a single ODE/SDE solve — use "
            "pipekit_jax.truncated_scan at that layer."
        )
    raise ValueError(f"Unrecognised adjoint spec: {spec!r}")

truncated_scan(f, init, xs, *, k=1)

jax.lax.scan with gradients truncated to the trailing k steps.

The carry leaving every step before T - k is wrapped in stop_gradient, so reverse-mode differentiation sees only the last k transitions through the carry. Each step's own output y_t still receives gradients with respect to that step's inputs — only the flow through history is cut, which is exactly the ROAD-EnKF / truncated-BPTT training regime. k >= T is the identity transformation (a plain scan).

Memory for the backward pass is O(k) in the rollout length instead of O(T), and for chaotic step functions the truncation acts as gradient regularisation (the exact adjoint norm grows like exp(lambda_max * T)).

Parameters:

Name Type Description Default
f Callable[[Carry, X], tuple[Carry, Y]]

Scan body (carry, x) -> (carry, y).

 required
init Carry

Initial carry.

 required
xs X

Stacked per-step inputs with leading axis T.

 required
k int

Number of trailing steps that propagate gradients through the carry. Must be at least 1.

 1

Returns:

Type Description
Carry

(final_carry, ys) exactly as jax.lax.scan would —
Y

the forward values are identical; only gradients differ.

Raises:

Type Description
ValueError

if k < 1 or xs has no leading axis.
Source code in packages/pipekit-jax/src/pipekit_jax/adjoints.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def truncated_scan[Carry, X, Y](
    f: Callable[[Carry, X], tuple[Carry, Y]],
    init: Carry,
    xs: X,
    *,
    k: int = 1,
) -> tuple[Carry, Y]:
    """``jax.lax.scan`` with gradients truncated to the trailing ``k`` steps.

    The carry leaving every step before ``T - k`` is wrapped in
    ``stop_gradient``, so reverse-mode differentiation sees only the
    last ``k`` transitions *through the carry*. Each step's own output
    ``y_t`` still receives gradients with respect to that step's
    inputs — only the flow through history is cut, which is exactly
    the ROAD-EnKF / truncated-BPTT training regime. ``k >= T`` is the
    identity transformation (a plain scan).

    Memory for the backward pass is O(k) in the rollout length instead
    of O(T), and for chaotic step functions the truncation acts as
    gradient regularisation (the exact adjoint norm grows like
    ``exp(lambda_max * T)``).

    Args:
        f: Scan body ``(carry, x) -> (carry, y)``.
        init: Initial carry.
        xs: Stacked per-step inputs with leading axis ``T``.
        k: Number of trailing steps that propagate gradients
            through the carry. Must be at least 1.

    Returns:
        ``(final_carry, ys)`` exactly as ``jax.lax.scan`` would —
        the forward values are identical; only gradients differ.

    Raises:
        ValueError: if ``k < 1`` or ``xs`` has no leading axis.
    """
    if k < 1:
        raise ValueError(f"truncated_scan needs k >= 1; got k={k}.")
    leaves = jax.tree.leaves(xs)
    if not leaves:
        raise ValueError("truncated_scan needs at least one stacked input in xs.")
    length = leaves[0].shape[0]
    cutoff = length - k

    def body(carry: Carry, indexed: tuple[Any, X]) -> tuple[Carry, Y]:
        i, x = indexed
        new_carry, y = f(carry, x)
        # jnp.where keeps the graph differentiable in the trailing
        # window and routes earlier steps through stop_gradient.
        cut = i < cutoff
        new_carry = jax.tree.map(
            lambda c: jnp.where(cut, jax.lax.stop_gradient(c), c), new_carry
        )
        return new_carry, y

    return jax.lax.scan(body, init, (jnp.arange(length), xs))

DiffraxForwardModel

A pipekit_cycle.ForwardModel backed by diffrax.diffeqsolve — ODE or SDE — that also satisfies SupportsAdjoint, so the gradient strategy of the forecast step is configuration. Requires the pipekit-jax[diffrax] extra.

pipekit_jax.diffrax_model

DiffraxForwardModel — a pipekit ForwardModel backed by diffrax.

Wraps diffrax.diffeqsolve so a vector field (ODE) or drift + diffusion pair (SDE) becomes a component that satisfies both pipekit_cycle.ForwardModel (step(state, dt) plus dt / state_signature) and pipekit_cycle.SupportsAdjoint structurally. The gradient strategy for the solve is declared with a spec from pipekit_cycle.adjoints and interpreted by pipekit_jax.to_diffrax_adjoint — so swapping, say, checkpointed for backsolve adjoints is configuration, not code.

Requires the pipekit-jax[diffrax] extra.

DiffraxForwardModel

Bases: Module

Advance a state with diffrax.diffeqsolve.

ODE by default; supply diffusion (and sde_key) for an SDE, in which case the Brownian path is a VirtualBrownianTree over each integration span. The adjoint spec selects how reverse-mode gradients flow through the solve (default: treeverse checkpointing, the robust choice — see pipekit_cycle.adjoints for the menu and the warning on BacksolveAdjoint).

pipekit cycles call step(state, dt) with no absolute time, so integration always runs over [0, dt] (autonomous dynamics); as_dynamics() exposes the explicit (state, t0, t1) form used by algorithm libraries (e.g. filterax AbstractDynamics).

Attributes:

Name Type Description
vector_field Callable[[Any, Any, Any], Any]

Drift f(t, y, args) -> dy/dt in diffrax convention.
solver AbstractSolver

Any diffrax solver (Tsit5() for non-stiff ODEs, Kvaerno5() for stiff, Euler()/Heun() for SDEs, ...).
dt0 float

Initial/fixed integrator step.
dt float

Default window length advertised to pipekit cycles.
adjoint Any

Adjoint spec (or a raw diffrax.AbstractAdjoint).
stepsize_controller AbstractStepSizeController

diffrax step-size controller; defaults to constant steps.
diffusion Callable[[Any, Any, Any], Any] | None

Optional g(t, y, args) turning the solve into an SDE dy = f dt + g dW. Per diffrax >= 0.7 semantics the return value must be a linear map from the Brownian increment into state space — e.g. lineax.DiagonalLinearOperator(sigma) for diagonal noise, or a (N_x, N_w) matrix.
brownian_shape tuple[int, ...] | None

Shape of the Brownian increment (SDE only). None derives it from the state — correct for diagonal noise. For non-square noise (a diffusion returning an (N_x, N_w) linear map with N_w != N_x) pass (N_w,) explicitly.
sde_key Array | None

PRNG key for the Brownian path (SDE only). The key is fixed at construction; refresh it per cycle through your DA state (e.g. DAState.extras["rng_key"]) by rebuilding with dataclasses.replace if independent paths per window are required.
max_steps int

Bound on integrator steps per solve.
Source code in packages/pipekit-jax/src/pipekit_jax/diffrax_model.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
class DiffraxForwardModel(eqx.Module):
    """Advance a state with ``diffrax.diffeqsolve``.

    ODE by default; supply ``diffusion`` (and ``sde_key``) for an SDE,
    in which case the Brownian path is a ``VirtualBrownianTree`` over
    each integration span. The adjoint spec selects how reverse-mode
    gradients flow through the solve (default: treeverse
    checkpointing, the robust choice — see
    ``pipekit_cycle.adjoints`` for the menu and the warning on
    ``BacksolveAdjoint``).

    pipekit cycles call ``step(state, dt)`` with no absolute time, so
    integration always runs over ``[0, dt]`` (autonomous dynamics);
    ``as_dynamics()`` exposes the explicit ``(state, t0, t1)`` form
    used by algorithm libraries (e.g. filterax ``AbstractDynamics``).

    Attributes:
        vector_field: Drift ``f(t, y, args) -> dy/dt`` in diffrax
            convention.
        solver: Any diffrax solver (``Tsit5()`` for non-stiff ODEs,
            ``Kvaerno5()`` for stiff, ``Euler()``/``Heun()`` for
            SDEs, ...).
        dt0: Initial/fixed integrator step.
        dt: Default window length advertised to pipekit cycles.
        adjoint: Adjoint spec (or a raw ``diffrax.AbstractAdjoint``).
        stepsize_controller: diffrax step-size controller; defaults to
            constant steps.
        diffusion: Optional ``g(t, y, args)`` turning the solve into
            an SDE ``dy = f dt + g dW``. Per diffrax >= 0.7 semantics
            the return value must be a *linear map* from the Brownian
            increment into state space — e.g.
            ``lineax.DiagonalLinearOperator(sigma)`` for diagonal
            noise, or a ``(N_x, N_w)`` matrix.
        brownian_shape: Shape of the Brownian increment (SDE only).
            ``None`` derives it from the state — correct for diagonal
            noise. For non-square noise (a diffusion returning an
            ``(N_x, N_w)`` linear map with ``N_w != N_x``) pass
            ``(N_w,)`` explicitly.
        sde_key: PRNG key for the Brownian path (SDE only). The key is
            fixed at construction; refresh it per cycle through your
            DA state (e.g. ``DAState.extras["rng_key"]``) by rebuilding
            with ``dataclasses.replace`` if independent paths per
            window are required.
        max_steps: Bound on integrator steps per solve.
    """

    vector_field: Callable[[Any, Any, Any], Any]
    solver: diffrax.AbstractSolver = diffrax.Tsit5()
    dt0: float = 0.01
    dt: float = 1.0
    adjoint: Any = eqx.field(static=True, default_factory=RecursiveCheckpointAdjoint)
    stepsize_controller: diffrax.AbstractStepSizeController = diffrax.ConstantStepSize()
    diffusion: Callable[[Any, Any, Any], Any] | None = None
    sde_key: jax.Array | None = None
    brownian_shape: tuple[int, ...] | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=4096)

    def __check_init__(self) -> None:
        if self.diffusion is not None and self.sde_key is None:
            raise ValueError(
                "SDE dynamics (diffusion is not None) need an sde_key for the "
                "VirtualBrownianTree."
            )

    def _terms(self, t0: Any, t1: Any, state: Any) -> Any:
        drift = diffrax.ODETerm(self.vector_field)
        if self.diffusion is None:
            return drift
        bm_shape = (
            self.brownian_shape if self.brownian_shape is not None else jnp.shape(state)
        )
        brownian = diffrax.VirtualBrownianTree(
            t0,
            t1,
            tol=self.dt0 / 2,
            shape=jax.ShapeDtypeStruct(bm_shape, jnp.result_type(state)),
            key=self.sde_key,
        )
        return diffrax.MultiTerm(drift, diffrax.ControlTerm(self.diffusion, brownian))

    def _solve(self, state: Any, t0: Any, t1: Any) -> Any:
        sol = diffrax.diffeqsolve(
            self._terms(t0, t1, state),
            self.solver,
            t0=t0,
            t1=t1,
            dt0=self.dt0,
            y0=state,
            stepsize_controller=self.stepsize_controller,
            adjoint=to_diffrax_adjoint(self.adjoint),
            max_steps=self.max_steps,
            saveat=diffrax.SaveAt(t1=True),
        )
        return jax.tree.map(lambda y: y[-1], sol.ys)

    def step(self, state: Any, dt: Any) -> Any:
        """Advance ``state`` by ``dt`` (autonomous: integrates ``[0, dt]``)."""
        return self._solve(state, jnp.zeros_like(jnp.asarray(dt)), dt)

    @property
    def state_signature(self) -> None:
        """No shape contract advertised; pipekit treats ``None`` as opt-out."""
        return None

    def with_adjoint(self, spec: Any) -> DiffraxForwardModel:
        """Return a copy configured with another adjoint spec."""
        return dataclasses.replace(self, adjoint=spec)

    def as_dynamics(self) -> Callable[[Any, Any, Any], Any]:
        """Return the ``(state, t0, t1) -> state`` callable form.

        Plugs directly into APIs that take explicit integration spans,
        e.g. filterax ``AbstractDynamics`` slots or vardax 4D-Var
        ``forward`` arguments.
        """
        return self._solve

state_signature property

No shape contract advertised; pipekit treats None as opt-out.

as_dynamics()

Return the (state, t0, t1) -> state callable form.

Plugs directly into APIs that take explicit integration spans, e.g. filterax AbstractDynamics slots or vardax 4D-Var forward arguments.

Source code in packages/pipekit-jax/src/pipekit_jax/diffrax_model.py 
138
139
140
141
142
143
144
145
def as_dynamics(self) -> Callable[[Any, Any, Any], Any]:
    """Return the ``(state, t0, t1) -> state`` callable form.

    Plugs directly into APIs that take explicit integration spans,
    e.g. filterax ``AbstractDynamics`` slots or vardax 4D-Var
    ``forward`` arguments.
    """
    return self._solve

step(state, dt)

Advance state by dt (autonomous: integrates [0, dt]).

Source code in packages/pipekit-jax/src/pipekit_jax/diffrax_model.py 
125
126
127
def step(self, state: Any, dt: Any) -> Any:
    """Advance ``state`` by ``dt`` (autonomous: integrates ``[0, dt]``)."""
    return self._solve(state, jnp.zeros_like(jnp.asarray(dt)), dt)

with_adjoint(spec)

Return a copy configured with another adjoint spec.

Source code in packages/pipekit-jax/src/pipekit_jax/diffrax_model.py 
134
135
136
def with_adjoint(self, spec: Any) -> DiffraxForwardModel:
    """Return a copy configured with another adjoint spec."""
    return dataclasses.replace(self, adjoint=spec)