def apply_gradients(self, *, grads: Any, **kwargs: Any) -> "TrainState":
    """Apply gradients exactly as Flax does and register the transition."""

    if not is_tracing():
        return super().apply_gradients(grads=grads, **kwargs)

    transaction = Transaction(
        batch=AUTOMATIC_BATCH,
        params=meta.unbox(self.params),
        opt_state=self._zkdv_optimizer_state(),
        overlap=self.zkdv_overlap,
        _proof=self.proof,
    )
    overwrite = isinstance(grads, Mapping) and OVERWRITE_WITH_GRADIENT in grads
    if overwrite:
        gradients = grads["params"]
        parameters = self.params["params"]
    else:
        gradients = grads
        parameters = self.params

    updates, optimizer = self.tx.update(gradients, self.opt_state, parameters)
    if overwrite:
        committed_updates = {
            "params": updates,
            OVERWRITE_WITH_GRADIENT: (
                grads[OVERWRITE_WITH_GRADIENT]
                - self.params[OVERWRITE_WITH_GRADIENT]
            ),
        }
    else:
        committed_updates = updates
    if transaction._attested:
        state = self.replace(step=self.step + 1, opt_state=optimizer, **kwargs)
        transaction.update(
            deltas=meta.unbox(committed_updates),
            params=meta.unbox(self.params),
            opt_state=state._zkdv_optimizer_state(),
        )
        return state

    updated_parameters = optax.apply_updates(parameters, updates)
    params = (
        {
            "params": updated_parameters,
            OVERWRITE_WITH_GRADIENT: grads[OVERWRITE_WITH_GRADIENT],
        }
        if overwrite
        else updated_parameters
    )
    state = self.replace(
        step=self.step + 1,
        params=params,
        opt_state=optimizer,
        **kwargs,
    )
    transaction.update(
        deltas=meta.unbox(committed_updates),
        params=meta.unbox(params),
        opt_state=state._zkdv_optimizer_state(),
    )
    return state

def lower(self, *args: Any, **kwargs: Any) -> Any:
    """Lower the fused executable without opening a transcript transaction."""

    values = self._signature.bind(args, kwargs)
    self._prepare(values)
    compiled_values, _ = self._execution(values)
    return self._compiled.lower(*compiled_values)

def accept(self, index: int, params: Any, optimizer: Any) -> None:
    """Record acceptance without copying the replay result back to host."""
    del params, optimizer
    with self._lock:
        if index not in self._active:
            raise RuntimeError("ZKDV accepted an inactive replay snapshot")

def activate(self, index: int, params: Any, optimizer: Any) -> None:
    """Open a transaction and asynchronously copy its pre-state when due."""
    wait_for = None
    while True:
        with self._lock:
            self._reap_locked()
            if index in self._active:
                raise RuntimeError(
                    f"ZKDV replay snapshot {index} is already active"
                )
            due = index not in self._snapshots and (
                not self._snapshots or index % self.snapshot_interval == 0
            )
            if not due or self._free:
                self._active.add(index)
                break
            wait_for = self._retired[0] if self._retired else None
        if wait_for is None:
            raise RuntimeError("ZKDV replay snapshot pool is exhausted")
        # If copies cannot keep up with the selected interval, expose real
        # backpressure rather than silently violating the replay bound.
        wait_for.result()

    if not due:
        return
    with self._lock:
        destination = self._free.pop()
        dispatching = Event()
        future = self._executor.submit(
            self._copy_to_host,
            params,
            optimizer,
            destination,
            dispatching,
        )
        self._snapshots[index] = _Snapshot(index, future)
        self._prune_locked()
    # Wait only until the copy worker owns the GIL immediately before the
    # compiled call. Its input handler therefore registers source usage
    # before this thread can enter the next donated JAX call. This is not a
    # transfer wait: the host-valued result remains on ``future``.
    dispatching.wait()

def close(self) -> None:
    """Release the dispatcher after the owning pipeline has joined."""
    self.join()
    self._executor.shutdown(wait=True, cancel_futures=False)

def join(self) -> None:
    """Drain copies at an explicit pipeline synchronization boundary."""
    with self._lock:
        futures = tuple(snapshot.future for snapshot in self._snapshots.values())
        futures += tuple(self._retired)
    for future in futures:
        future.result()
    with self._lock:
        self._reap_locked()

def pre_state(self, index: int) -> tuple[Any, Any]:
    """Restore the newest eligible snapshot and replay only its short suffix."""
    with self._lock:
        eligible = [current for current in self._snapshots if current <= index]
        if not eligible:
            raise RuntimeError("ZKDV has no replay snapshot for sampled update")
        start = max(eligible)
        snapshot = self._snapshots[start]
        batches = []
        for current in range(start, index):
            if current not in self._batches:
                raise RuntimeError(f"ZKDV checkpoint is missing batch {current}")
            batches.append(self._batches[current])

    # This is intentionally the first wait on a retained snapshot, and
    # this method is called only after Rust selects an update for checking.
    host_params, host_optimizer = snapshot.future.result()
    with jax.profiler.TraceAnnotation("zkdv.snapshot.restore"):
        params, optimizer = self._restore(host_params, host_optimizer)
    for packed_batch in batches:
        batch = jax.device_put(
            jax.device_get(packed_batch),
            self._batch_sharding,
        )
        params, optimizer = self._advance(params, optimizer, batch)
    jax.block_until_ready((params, optimizer))
    return params, optimizer

def record(self, index: int, batch: Any) -> None:
    """Retain one independently owned host batch for a possible suffix."""
    with self._lock:
        if index in self._batches:
            raise RuntimeError(f"ZKDV checkpoint batch {index} is out of order")
        self._batches[index] = batch
        self._prune_locked()

def release(self, index: int) -> None:
    """Release one transaction; never fence a discarded host transfer."""
    with self._lock:
        self._active.discard(index)
        self._prune_locked()

def pack(
    self, values: tuple[Any, ...], host_values: tuple[Any, ...] | None = None
) -> np.ndarray:
    """Encode dynamic call inputs without adding work to the JIT executable."""

    leaves = jax.tree.leaves(values)
    if host_values is None:
        host_values = tuple(leaves[position] for position in self.sources)
    host = jax.device_get(host_values)
    encoded = []
    for value, spec in zip(host, self.specs, strict=True):
        if spec.key:
            value = jax.device_get(jax.random.key_data(value))
        array = np.ascontiguousarray(np.asarray(value, dtype=spec.dtype))
        encoded.append(array.view(np.uint8).reshape(-1))
    raw = np.concatenate(encoded)
    padding = (-raw.size) % np.dtype(np.uint32).itemsize
    if padding:
        raw = np.pad(raw, (0, padding))
    return np.ascontiguousarray(raw.view(np.uint32).reshape(1, -1))

def program(self, stage: Callable[..., tuple[Any, Any]]) -> Callable[..., Any]:
    """Expose the full transition used by checkpoint replay and checks."""

    def attested(params, opt_state, batch):
        values = self.rebuild(params, opt_state, batch)
        return stage(*values)

    return attested

def training_stage(
    self, stage: Callable[..., tuple[Any, ...]]
) -> Callable[..., tuple[Any, ...]]:
    """Present parameter buffers separately so only they may be donated."""

    def lowered(params, opt_state, *inputs):
        return stage(*self.rebuild_from_inputs(params, opt_state, inputs))

    return lowered