Train a Flax model with TrainState¶

Use zkdv.jax.contrib.TrainState when your JAX training loop already follows Flax's TrainState pattern. It behaves like Flax's state, but also declares the optimizer transition that DVBench needs to replay sampled updates.

1. Install the JAX dependencies¶

pip install "dvbench[jax]"

This installs DVBench together with ZKDV's JAX, Flax, and Optax integration.

2. Start a JAX proof¶

Connect to DVBench, register the experiment, and start its proof before you create the training state.

from dvbench import Client

client = Client(API_URL)
experiment = client.experiment("your-name/jax-train-state")
proof = experiment.proof(backend="jax")
proof.start(rolling_window=4, probe_ratio=0.01)

The object returned by experiment.proof() is the full zkdv.jax.ZKDV driver with DVBench upload and cleanup behavior added. Its start(), jit(), and other ZKDV methods remain available; call finish() at the end to finalize and upload the proof.

import zkdv

config = zkdv.ZKDVConfig(
    sampling_policy=zkdv.SamplingPolicy(
        update_check_probability=0.01,
    ),
)
proof = experiment.proof(backend="jax", config=config)

3. Create the model and state¶

Initialize the model normally, then replace flax.training.train_state.TrainState with zkdv.jax.contrib.TrainState. The familiar apply_fn, params, and tx arguments are unchanged; proof identifies the ZKDV driver that owns this state's updates.

import jax
import jax.numpy as jnp
import optax
from flax import linen as nn
import zkdv


class TinyClassifier(nn.Module):
    @nn.compact
    def __call__(self, tokens):
        return nn.Dense(2)(tokens.astype(jnp.float32))


model = TinyClassifier()
params = model.init(
    jax.random.key(0),
    jnp.ones((1, 4), dtype=jnp.uint32),
)["params"]

state = zkdv.jax.contrib.TrainState.create(
    proof=proof,
    apply_fn=model.apply,
    params=params,
    tx=optax.adam(1e-2),
)

The wrapper inherits Flax's complete state API:

• create() initializes step and opt_state from the supplied Optax transformation.
• apply_gradients(grads=..., **updates) applies the Optax update, increments step, and replaces any additional fields passed in updates.
• replace(**updates) returns a state with selected fields replaced.
• Applications may subclass it to add Linen state or other PyTree fields.

It adds these three fields:

All standard Flax fields remain available: step, apply_fn, params, tx, and opt_state. Flax's OVERWRITE_WITH_GRADIENT convention for FP8 parameters is also handled by the wrapper.

4. Compile the training step¶

Replace jax.jit with proof.jit at the function that calls state.apply_gradients(). The model, loss, gradient calculation, and return value can otherwise remain ordinary JAX code.

@proof.jit(donate_argnums=(0,))
def train_step(state, batch):
    tokens = batch[:, :4]
    labels = batch[:, 4].astype(jnp.int32)

    def loss(params):
        logits = state.apply_fn({"params": params}, tokens)
        return optax.softmax_cross_entropy_with_integer_labels(
            logits,
            labels,
        ).mean()

    grads = jax.grad(loss)(state.params)
    return state.apply_gradients(grads=grads)

proof.jit accepts the usual jax.jit options, including in_shardings, out_shardings, static arguments, donation, device, backend, inline, and compiler options. Donating the state argument is supported; ZKDV retains the pre-update values required by a sampled check.

The wrapper automatically treats every dynamic input that is not part of the state as committed batch evidence. Pass data-dependent randomness, such as a PRNG key, as another dynamic argument:

@proof.jit
def train_step(state, batch, key):
    augmented = augment(batch, key)
    grads = jax.grad(loss)(state.params, augmented)
    return state.apply_gradients(grads=grads)

5. Train normally¶

batch = jnp.asarray(
    [
        [0, 0, 0, 1, 0],
        [0, 1, 0, 1, 0],
        [1, 0, 1, 0, 1],
        [1, 1, 1, 0, 1],
    ],
    dtype=jnp.uint32,
)

for _ in range(10):
    state = train_step(state, batch)

The state also works with jax.device_put and ordinary jax.jit. An ordinary jax.jit call has normal Flax behavior, but only a step compiled with proof.jit contributes an update to the DVBench proof.

6. Finish and save the checkpoint¶

Wait for all commitments and sampled checks before saving the parameters used for evaluation.

from pathlib import Path
from flax import serialization

proof.finish()
Path("jax-checkpoint.msgpack").write_bytes(
    serialization.to_bytes(state.params),
)
client.close()

finish() also uploads the completed transcript and releases temporary proof storage. It is safe to call more than once.

Carry additional Flax state¶

Subclass the ZKDV wrapper just as you would subclass Flax TrainState. Pass new field values to create() and updated values to apply_gradients().

from typing import Any
from flax import struct


class ModelTrainState(zkdv.jax.contrib.TrainState):
    batch_stats: Any = struct.field(pytree_node=True)


state = ModelTrainState.create(
    proof=proof,
    apply_fn=model.apply,
    params=params,
    tx=optax.adam(1e-2),
    batch_stats=initial_batch_stats,
)


@proof.jit
def train_step(state, batch):
    grads, next_batch_stats = compute_gradients(state, batch)
    return state.apply_gradients(
        grads=grads,
        batch_stats=next_batch_stats,
    )

Every dynamic state leaf other than params is included in optimizer-state continuity. This includes step, Optax state, and dynamic fields added by a subclass.

Supply an exportable replay function¶

The automatically generated sampled program must pass jax.export. If the production apply_fn uses a custom kernel or another non-exportable operation, provide an equivalent portable implementation through attest_apply_fn:

state = zkdv.jax.contrib.TrainState.create(
    proof=proof,
    apply_fn=fast_model_apply,
    attest_apply_fn=portable_model_apply,
    params=params,
    tx=optimizer,
)

ZKDV substitutes portable_model_apply only during sampled replay. It must have the same call signature and produce the same values used by the loss; a mismatch causes verification to fail.

Choose an overlap policy¶

Set zkdv_overlap when creating or replacing the state:

state = zkdv.jax.contrib.TrainState.create(
    proof=proof,
    apply_fn=model.apply,
    params=params,
    tx=optimizer,
    zkdv_overlap=zkdv.overlap.NEVER,
)

• zkdv.overlap.UNCHECKED is the default. Ordinary commitments pipeline, but a sampled check acts as a fence.
• zkdv.overlap.ALWAYS allows training and a sampled verification to overlap, using more live device-buffer generations.
• zkdv.overlap.NEVER waits for the current commitment and any sampled check before admitting the next update. It uses the least concurrent residency.

Start with UNCHECKED; choose ALWAYS when throughput matters and memory can accommodate the extra live buffers, or NEVER when minimizing residency is more important.

Common integration errors¶

• No transcript update: the training step still uses jax.jit. Decorate the function containing apply_gradients() with proof.jit.
• Proof identity error: the state was created with a different proof object. Use the exact instance returned for this experiment.
• Attested program already registered: remove the manual proof.attest function when using the TrainState integration.
• Static argument error: automatic attestation fixes static arguments on the first training invocation. Keep their values unchanged for that compiled step, or make the value a supported dynamic input.
• Export or replay failure: pass randomness explicitly and use attest_apply_fn only for a deterministic, equivalent implementation.

A complete working program is available in examples/jax/proof_trainstate.py.