API¶

Everything below is importable from pytest_torch directly, e.g. from pytest_torch import assert_export_parity.

Generated suites¶

ModelSuite ¶

Base class for a generated model test suite (it is not collected itself).

Override :meth:make_model and :meth:make_inputs (and :meth:loss for the "training" checks), and pick the checks with checks; see CHECKS for all names. check_options passes keyword arguments to the assertion behind a test, keyed by the test name without test_, e.g. {"inference_contract": {"skip": ["device_parity"]}, "overfits_one_batch": {"steps": 500}}. Any generated test can be overridden like a normal method.

make_model ¶

make_model()

Return a fresh model. Weights are seeded by torch_seed.

make_inputs ¶

make_inputs(make_tensor) -> Any

Return example inputs (tensor, tuple of args or dict of kwargs) on the test's device.

loss ¶

loss(output, inputs)

Return a scalar training loss for output; needed by the "training" checks.

frozen_parameters ¶

frozen_parameters() -> Sequence[str]

Patterns for the "frozen" check; override to depend on self.config (an empty result skips the check for that config).

load_config ¶

load_config(source: Any) -> Any

Turn an entry of configs into a config: YAML/JSON files are parsed (YAML needs PyYAML), anything else is returned as is. Override for Hydra/OmegaConf composition.

config_id `classmethod` ¶

config_id(source: Any) -> str

Test id for a config: the file name without extension, or config<i>.

Assertions¶

Tensor/pytree assertions with readable failure messages.

assert_tensors_close ¶

assert_tensors_close(actual: Tensor, expected: Tensor, *, rtol: float | None = None, atol: float | None = None, check_dtype: bool = True, check_device: bool = True, equal_nan: bool = True) -> None

Assert two tensors are close, reporting where and how badly they differ.

Floating-point tensors use dtype-aware default tolerances; complex tensors compare the complex difference (real and imaginary parts); integer and bool tensors compare exactly unless rtol/atol is given. NaNs must sit at the same positions, and count as equal unless equal_nan=False.

assert_outputs_close ¶

assert_outputs_close(actual: Any, expected: Any, *, rtol: float | None = None, atol: float | None = None, check_dtype: bool = True, check_device: bool = True, equal_nan: bool = True) -> None

Pytree-aware assert_tensors_close for dicts, tuples, dataclasses, ModelOutputs.

Reports the path of every mismatching entry (e.g. output['aux'][0]), missing and unexpected entries, and structure changes, including empty containers ([] vs {}). A tuple and a list with the same contents count as the same. Non-tensor leaves are compared with == (floats with a tolerance).

assert_finite ¶

assert_finite(obj: Any, name: str = 'tensor') -> None

Assert that every floating-point tensor in obj has no NaN/Inf values.

assert_batch_invariant ¶

assert_batch_invariant(model, inputs: Any, **tolerances) -> None

Assert model(batch)[i] ~ model(batch[i:i+1])[0] for every sample i.

inputs may be a tensor, a tuple of positional args, or a dict of kwargs (tensors whose leading dim equals the batch size are sliced). Outputs may be arbitrary pytrees. The model is expected to be in eval mode already.

Output specifications¶

Declarative output specifications: shapes, dtypes, finiteness, ranges, probabilities.

spec = {
    "logits": TensorSpec(("B", 10), dtype=torch.float32),
    "probs": Probabilities(("B", 10)),
    "hidden": TensorSpec(("B", "T", 64)),
}
assert_output_spec(model(x), spec, B=8)

Shape entries: an int (exact size), a str (a symbolic dim: every use must have the same size), None (any size) or one ... (any number of dims). Symbols are bound by keyword arguments or by their first occurrence. The spec mirrors the output's structure (dicts, lists/tuples, dataclasses, namedtuples, Hugging Face ModelOutputs).

TensorSpec `dataclass` ¶

What one output tensor must look like. Unset fields are not checked.

finite (default True) rejects NaN and Inf. min/max bound every value (inclusive). dtype may be a dtype or a tuple of allowed dtypes.

Probabilities `dataclass` ¶

Bases: TensorSpec

Probabilities along dim: values in [0, 1] summing to 1 (within atol).

With log=True: log-probabilities, values <= 0 whose logsumexp is 0.

assert_output_spec ¶

assert_output_spec(output: Any, spec: Any, **dims: int) -> dict[str, int]

Assert output matches spec and return the bound symbolic dims.

Keys present in the output but absent from the spec are ignored; None in the spec skips an entry. Every problem is reported, each with its output path.

Checkpoints¶

Checkpoints: saving and reloading preserves predictions; old checkpoints still load.

assert_checkpoint_roundtrip(make_model, inputs) # save -> fresh model -> load -> same outputs assert_checkpoint_loads(make_model, "ckpts/v1.pt") # yesterday's checkpoint, today's code

assert_checkpoint_roundtrip ¶

assert_checkpoint_roundtrip(make_model: Callable[[], Module], inputs: Any, *, model: Module | None = None, format: str = 'torch', rtol: float | None = None, atol: float | None = None) -> None

Save model's state_dict, load it into a fresh make_model() and compare outputs.

The fresh model is built under a different seed, so anything not restored from the checkpoint shows up as a difference. model defaults to make_model(); both are put in eval mode on the device of model. format is "torch" (torch.save) or "safetensors". On failure the message lists the usual culprits: tensors kept as plain attributes, non-persistent buffers, and tied weights that came apart.

assert_checkpoint_loads ¶

assert_checkpoint_loads(make_model: Callable[[], Module], checkpoint: str | PathLike | Mapping[str, Any], *, key: str | None = None, strip_prefix: str | None = None, strict: bool = True, weights_only: bool = True) -> torch.nn.Module

Load an existing checkpoint into make_model() and report every incompatibility.

checkpoint is a path or an already loaded object; the state_dict is found at key or under a common name (state_dict, model_state_dict, model, ...). strip_prefix removes e.g. "module." (DataParallel/DDP). With strict=False only shape mismatches fail. Returns the model with the weights loaded.

Stateful and streaming models¶

Stateful and streaming models: chunked processing, resets and request isolation.

report = assert_streaming_equivalent(
    full_runner=model.forward,
    streaming_runner=model.forward_chunk,
    inputs=sequence,
    chunk_sizes=[1, 7, 32],
    reset=model.reset_cache,
)

Three checks, reported separately:

chunked: for every chunk size, feeding the sequence chunk by chunk and concatenating the outputs agrees with one full pass (under the chosen equivalence);
reset: after reset() (or a fresh init_state()), a second streaming pass reproduces the first exactly;
isolation: one request does not contaminate another. With internal state: request A, reset, request B must equal B alone. With explicit state: A and B interleaved chunk by chunk, each with its own state, must each equal their isolated run.

The expected equivalence is user-defined, because some models intentionally behave differently when chunked (look-ahead, chunk-local attention, downsampling): see :class:StreamingEquivalence and the equivalence argument.

StreamingEquivalence `dataclass` ¶

How streamed outputs must relate to the full pass.

rtol/atol: tolerances (default: the streaming tolerance, configurable in [tool.pytest_torch.tolerance]). skip_first/skip_last: positions to ignore at the start/end of each chunk, e.g. skip_last=2 for a model that looks two steps ahead and so can't finish a chunk's last outputs without the next chunk. skip_warmup: positions to ignore at the start of the whole sequence.

StreamingReport `dataclass` ¶

assert_streaming_equivalent ¶

assert_streaming_equivalent(full_runner: Callable, streaming_runner: Callable, inputs: Any, *, chunk_sizes: Sequence[int] = (1, 7, 32), reset: Callable[[], Any] | None = None, init_state: Callable[[], Any] | None = None, equivalence: Equivalence | Mapping[Any, Equivalence] = None, seq_dim: int = 1, combine: Callable[[list], Any] | None = None, other_inputs: Any = None, checks: Sequence[str] = ('chunked', 'reset', 'isolation'), raise_on_failure: bool = True) -> StreamingReport

Check a stateful model's streaming path against its full path. See the module docs.

full_runner(inputs) and streaming_runner(chunk) are called like a model (a tensor, a tuple of positional arguments, or a dict of keyword arguments). With init_state, the streaming runner is called as streaming_runner(chunk, state) and returns (output, new_state). Inputs are split along seq_dim; tensors of another length (embeddings, flags) go to every chunk whole; the sequence length is the longest size along seq_dim.

equivalence: None (outputs agree within the streaming tolerance), a :class:StreamingEquivalence (tolerances, positions to skip), a callable (full_output, streamed_output, chunk_size) that raises AssertionError or returns False on a mismatch, or a dict of those keyed by chunk size ("default" for the rest). other_inputs is the second request for the isolation check (default: the inputs reversed in time). Returns a :class:StreamingReport.

Dataset evaluation¶

Dataset evaluation: run a loader through a model, aggregate metrics, guard regressions.

def accuracy(output, batch):
    return (output.argmax(-1) == batch[1]).float()        # per sample: enables paired tests

def test_accuracy(model, val_loader, metric_baseline):
    metric_baseline(model, val_loader, {"accuracy": accuracy},
                    context={"dataset": "sst2", "split": "validation", "preprocessing": "v3"},
                    minimum={"accuracy": 0.9})

Record baselines with pytest --torch-update-metrics. A baseline remembers what it measured (context, input fingerprint, sample selection, metric configuration, checkpoint), and a later run is only compared with a baseline that measured the same thing. Metrics with per-sample values are compared paired, with a bootstrap confidence interval, so a small drop on a small dataset isn't mistaken for a regression.

Evaluation `dataclass` ¶

Result of :func:evaluate_detailed.

PairedDifference `dataclass` ¶

Paired change of one metric versus its baseline: drop > 0 means worse.

MetricResult `dataclass` ¶

evaluate ¶

evaluate(model, loader: Iterable, metrics: Mapping[str, Any], *, inputs: Callable[[Any], Any] = default_inputs, max_batches: int | None = None) -> dict[str, float]

Run loader through model (eval mode, inference mode, model's device).

A metric is fn(output, batch) returning a scalar (averaged over the dataset weighted by batch size) or one value per sample (averaged over samples), or an object with update(output, batch) and compute() for metrics that need every prediction (F1, AUC, ...). inputs extracts the model inputs from a batch. max_batches caps the work, e.g. in CI.

evaluate_detailed ¶

evaluate_detailed(model, loader: Iterable, metrics: Mapping[str, Any], *, inputs: Callable[[Any], Any] = default_inputs, max_batches: int | None = None, context: Mapping[str, Any] | None = None, fingerprint_model: bool = True) -> Evaluation

Like :func:evaluate, plus per-sample values and provenance.

A metric fn(output, batch) returning a 1-D tensor with one value per sample is per-sample (enables paired comparisons); a scalar is a batch mean, weighted by batch size. Objects with update(output, batch) / compute() cover metrics that need every prediction; give them a config attribute (averaging, thresholds, ignored labels) so baselines record it. context describes the evaluation (dataset, split, preprocessing version, checkpoint label).

paired_difference ¶

paired_difference(current: Tensor, baseline: Tensor, *, higher_is_better: bool = True, confidence: float = 0.95, resamples: int = 2000, seed: int = 0) -> PairedDifference

Mean per-sample drop (> 0: worse) with a seeded percentile-bootstrap interval.

check_metrics ¶

check_metrics(store, key: str, values: Mapping[str, float] | Evaluation, *, update: bool = False, strict: bool = False, tolerance: float | Mapping[str, float] = 0.01, higher_is_better: bool | Mapping[str, bool] = True, minimum: Mapping[str, float] | None = None, maximum: Mapping[str, float] | None = None, confidence: float = 0.95, resamples: int = 2000) -> MetricResult

Compare an evaluation with the stored baseline and with absolute thresholds.

The baseline must have measured the same thing (context except checkpoint, sample count and selection, input fingerprint, metric configuration); otherwise the check fails and says what changed. Then, per metric:

with per-sample values on both sides: a paired comparison. A regression is a mean drop larger than tolerance whose bootstrap confidence interval (confidence) excludes zero, so noise on a small dataset doesn't fail the test;
otherwise: a regression is a drop larger than tolerance (absolute).

Improvements never fail. minimum/maximum apply on every run. A missing baseline skips the test unless thresholds were given (they still apply) or strict is set.

comparability_problems ¶

comparability_problems(baseline: Mapping[str, Any], current: Mapping[str, Any]) -> list[str]

Why current measured something other than baseline (empty: comparable).

model_fingerprint ¶

model_fingerprint(model) -> str | None

Short hash of a model's parameters and buffers (what was evaluated).

Failure reproduction¶

Failure reproduction bundles: record what a failing test saw, replay it anywhere.

pytest --torch-repro-dir=repro                 # every failing test writes repro/<test>.pt
pytest --torch-repro=repro/<test>.pt           # rerun just that test, same seed and inputs
pytest --torch-repro=repro/<test>.pt --torch-device=cpu   # ...on another device

A bundle holds the test id and parameters, seed, device, every tensor make_tensor produced (replayed bit for bit, so platform RNG differences don't matter), the values passed to :func:record_repro, the active tolerances and the environment. Model weights are only stored up to --torch-repro-max-weights-mb (default 0: off), since models built under torch_seed are reproducible from the seed.

Torch-free at import: the plugin imports this module at startup.

record_repro ¶

record_repro(name: str, value: Any) -> Any

Store value in the failure bundle of the running test and return it.

When replaying a bundle, return the recorded value instead (for nn.Modules: the module, with recorded weights loaded into it if the bundle has them). Values may be tensors or nested dicts/lists/tuples of tensors and plain Python values. Outside --torch-repro-dir / --torch-repro runs this just returns value.

batch = record_repro("batch", next(iter(loader)))
model = record_repro("model", model)

load_repro ¶

load_repro(path: str | PathLike) -> dict

Load a bundle (e.g. to inspect bundle["inputs"] in a notebook).

Inference contract¶

One-call inference contract check.

report = check_inference_contract(model, example_inputs)

runs a battery of checks that every inference model should satisfy and reports which failed and why it likely failed.

ContractReport `dataclass` ¶

CheckResult `dataclass` ¶

check_inference_contract ¶

check_inference_contract(model: Callable, inputs: Any, *, checks: Sequence[str] | None = None, skip: Sequence[str] = (), rtol: float | None = None, atol: float | None = None, devices: Sequence[str] | None = None, device_rtol: float | None = None, device_atol: float | None = None, dtypes: Sequence[dtype] = (), dtype_rtol: float | None = None, dtype_atol: float | None = None, raise_on_failure: bool = True) -> ContractReport

Run a battery of inference checks and return a :class:ContractReport.

inputs may be a tensor, a tuple of positional args, or a dict of kwargs; they are moved to the model's device automatically. Outputs may be any pytree.

Checks (see CHECKS): eval mode, finite outputs, determinism, batch invariance, device parity (vs devices; default: the other available devices), dtype parity (only for dtypes, e.g. (torch.bfloat16,)), and no parameter/buffer mutation. Use checks= / skip= to select; selecting no check at all is a ValueError, raised before the model runs. Tolerances: rtol/atol for same-device checks (dtype defaults when None), device_* and dtype_* for the parity checks.

The first output is copied right away, so a model returning a buffer it overwrites on later calls can't pass the determinism check by aliasing. no_state_mutation covers every registered parameter and buffer, non-persistent buffers included, and reports state added or removed during inference; plain Python attributes (counters, caches stored as ordinary attributes) are outside its scope. eval_mode alone doesn't run the model.

Layer divergence¶

Find the first layer where two models (or two runs of one model) diverge.

report = localize_divergence(model_cpu, model_gpu, example_inputs)
print(report)

Forward hooks record every module's inputs and outputs in both models; layers are matched by name (override with module_map) and compared in execution order. Limitations: layers fused or skipped by torch.compile / TorchScript don't fire hooks, and later layers inherit earlier errors, so always start at the first row.

DivergenceReport `dataclass` ¶

LayerDiff `dataclass` ¶

localize_divergence ¶

localize_divergence(model_a, model_b, inputs: Any, *, inputs_b: Any = None, rtol: float | None = None, atol: float | None = None, module_map: Mapping[str, str] | None = None, leaf_only: bool = True) -> DivergenceReport

Compare per-layer activations of model_a (reference) and model_b.

inputs are moved to each model's device; floating inputs for model_b are cast to its parameter dtype (pass inputs_b to control this yourself). module_map maps names in model_a to names in model_b when they differ. leaf_only=False also records container modules.

Invariance and layouts¶

Invariance and metamorphic checks: padding, masks, permutations, custom transforms.

assert_invariant ¶

assert_invariant(model, inputs: Any, transform: Callable[[Any], Any], *, name: str = 'the transform', rtol: float | None = None, atol: float | None = None) -> None

Assert model(transform(x)) == model(x) (output is invariant to transform).

assert_equivariant ¶

assert_equivariant(model, inputs: Any, input_transform: Callable[[Any], Any], output_transform: Callable[[Any], Any], *, name: str = 'the transform', rtol: float | None = None, atol: float | None = None) -> None

Assert model(input_transform(x)) == output_transform(model(x)).

assert_permutation_invariant ¶

assert_permutation_invariant(model, inputs: Any, *, dim: int = 1, seed: int = 0, rtol=None, atol=None) -> None

Assert the output doesn't change when dim of the inputs is permuted (set models).

Every tensor whose size along dim equals that of the first input tensor is permuted consistently.

assert_permutation_equivariant ¶

assert_permutation_equivariant(model, inputs: Any, *, dim: int = 1, seed: int = 0, rtol=None, atol=None) -> None

Assert permuting dim of the inputs permutes the same dim of the outputs.

Output tensors are permuted only if their size along dim equals the permuted size, so pooled outputs must be invariant. A hidden size that coincidentally equals the sequence length would be misinterpreted.

assert_padding_invariant ¶

assert_padding_invariant(model, inputs: Mapping[str, Any], *, extra: int = 4, side: str = 'right', pad_value: int = 0, ids_key: str = 'input_ids', mask_key: str = 'attention_mask', seq_dim: int = 1, rtol: float | None = None, atol: float | None = None) -> None

Assert extra padding doesn't change the outputs for the real tokens.

Every input tensor whose seq_dim size equals the sequence length gets extra positions appended (side="right") or prepended ("left"): pad_value for ids/other tensors, 0 for the attention mask. Outputs that are sequence-aligned are compared on the original positions only; pooled outputs are compared as is. Models with absolute position ids can legitimately differ under left padding unless you pass position_ids explicitly; prefer side="right" for those.

assert_mask_invariant ¶

assert_mask_invariant(model, inputs: Mapping[str, Any], *, ids_key: str = 'input_ids', mask_key: str = 'attention_mask', vocab_size: int | None = None, seed: int = 0, rtol: float | None = None, atol: float | None = None) -> None

Assert that the content of masked-out positions doesn't affect the real tokens.

Replaces every masked id with a different random id and compares the outputs at the unmasked positions (sequence-aligned outputs, batch dim 0 and sequence dim 1) and all pooled outputs. Catches attention-mask bugs, e.g. padding leaking into pooling.

noncontiguous ¶

noncontiguous(t: Tensor) -> torch.Tensor

Same values and shape as t, but not contiguous in memory.

With 2+ dims the last two are stored column-major (like a transposed tensor), so .view() across them fails just as it would on real transposed data; 1-D tensors become a strided view.

assert_layout_invariant ¶

assert_layout_invariant(model, inputs: Any, *, layouts: tuple[str, ...] = ('noncontiguous', 'channels_last'), rtol: float | None = None, atol: float | None = None) -> None

Assert outputs don't depend on the memory layout of the inputs.

Each layout gets copies of the inputs with identical values: "noncontiguous" (strided views; catches .view() on non-contiguous tensors and kernels that assume contiguity) and "channels_last" (for 4-D inputs; skipped if there are none). Tolerance: the layout_invariance default, since kernels may legitimately differ by rounding between layouts.

Parity¶

Parity checks: eager vs deployment targets, models vs references, quantized vs fp32.

assert_compile_parity ¶

assert_compile_parity(model, inputs: Any, *, rtol: float | None = None, atol: float | None = None, runs: int = 2, **compile_kwargs) -> None

Assert torch.compile(model) matches eager mode.

The compiled model runs runs times (the second call exercises the cached graph and guards). Extra kwargs go to torch.compile, e.g. backend="eager" to test graph capture without needing a C++ compiler.

assert_torchscript_parity ¶

assert_torchscript_parity(model, inputs: Any, *, mode: str = 'trace', rtol: float | None = None, atol: float | None = None) -> None

Assert the TorchScript version (mode = "trace" or "script") matches eager.

inputs must be a tensor or a tuple of tensors (no kwargs). Note that tracing only records the control-flow path taken for these inputs.

assert_onnx_parity ¶

assert_onnx_parity(model, inputs: Any, *, rtol: float | None = None, atol: float | None = None, export_kwargs: Mapping[str, Any] | None = None, providers: tuple[str, ...] = ('CPUExecutionProvider',)) -> None

Export to ONNX, run with ONNX Runtime, and compare against PyTorch (on CPU).

inputs must be a tensor or a flat tuple of tensors. The ONNX exporter's defaults differ between torch versions: pin them via export_kwargs (e.g. {"dynamo": False, "opset_version": 17}). Needs pip install pytest-torch[onnx].

assert_export_parity ¶

assert_export_parity(model, inputs: Any, *more_inputs: Any, dynamic_batch: bool = False, dynamic_shapes: Any = None, strict: bool | None = None, rtol: float | None = None, atol: float | None = None)

Export with torch.export and assert the exported program matches eager.

more_inputs are extra input sets (e.g. another batch size) that the same exported program must accept and match on; that needs dynamic shapes: dynamic_batch=True marks dim 0 of every batch-aligned input as dynamic, or pass torch.export's own dynamic_shapes. strict is forwarded when set. Returns the ExportedProgram (e.g. to save it or lower it further).

assert_aoti_parity ¶

assert_aoti_parity(model, inputs: Any, *more_inputs: Any, dynamic_batch: bool = False, dynamic_shapes: Any = None, rtol: float | None = None, atol: float | None = None, inductor_configs: Mapping[str, Any] | None = None) -> None

Export, compile ahead-of-time with AOTInductor, load the package and compare to eager.

This is the path to running a model without Python (C++ servers, edge). Needs torch >= 2.6 and a C++ compiler; skips the test otherwise. Compilation takes a while (tens of seconds), so keep these tests few. Arguments as in :func:assert_export_parity.

assert_autocast_parity ¶

assert_autocast_parity(model, inputs: Any, *, dtype: dtype = torch.bfloat16, rtol: float | None = None, atol: float | None = None) -> None

Assert the model under torch.autocast(dtype) stays close to full precision.

This is how mixed precision usually runs (weights stay fp32, matmuls and convs run in dtype), unlike casting the whole model as the contract's dtype_parity check does. Outputs are compared in fp32. Skips if autocast isn't available for the model's device.

compare_outputs ¶

compare_outputs(reference: Any, candidate: Any, *, select=None) -> AgreementMetrics

Agreement metrics between two model outputs (default: first floating tensor).

assert_models_agree ¶

assert_models_agree(reference, candidate, inputs: Any, *, min_top1: float | None = None, min_cosine: float | None = None, max_abs: float | None = None, select=None, inputs_candidate: Any = None) -> AgreementMetrics

Assert candidate stays within an accuracy budget of reference.

Typical use: fp32 model vs. its quantized/pruned/distilled version. Budgets: min_top1 (fraction of identical argmax predictions), min_cosine (mean cosine similarity of output vectors) and max_abs (max absolute error). Set at least one. Returns the measured :class:AgreementMetrics so you can log them.

assert_matches_reference ¶

assert_matches_reference(model, reference: Any, inputs: Any, *, select: Callable | None = None, rtol: float | None = None, atol: float | None = None) -> None

Assert model(inputs) matches a reference implementation.

reference is either a model/callable (called with the same inputs) or a precomputed output (pytree). select extracts the part to compare, e.g. lambda out: out.logits; it is applied to model's output and to the output of a callable reference, but not to a precomputed reference.

assert_generation_equal ¶

assert_generation_equal(model, reference, inputs: Any, *, max_new_tokens: int = 16, generate: Callable | None = None, **generate_kwargs) -> None

Assert greedy decoding gives identical token sequences for model and reference.

inputs is an input_ids tensor or a dict of kwargs for generate. Extra kwargs (e.g. pad_token_id=0) go to model.generate. For non-Hugging-Face models pass generate=fn(model, inputs, max_new_tokens, kwargs) -> token tensor.

Training¶

Training sanity checks: gradients flow, are finite, are correct; the model can learn.

Every helper takes (model, inputs, loss_fn) where loss_fn(output) -> scalar tensor (close over your targets, e.g. lambda out: F.cross_entropy(out, y)). Helpers that take optimizer steps train a deep copy, so the model you pass is never modified.

StepRecord `dataclass` ¶

Telemetry of one training step, recorded by :func:run_training.

Per-optimizer values are keyed by optimizer name ("optimizer" for a single one). grad_norm, lrs and nonfinite_grads are read when optimizer.step() runs, so they describe the update that was applied.

loss `property` ¶

loss: float | None

The "loss" entry, or the first loss the step returned.

skipped `property` ¶

skipped: bool

No optimizer stepped (not called, or skipped by a GradScaler after an overflow).

TrainingRun `dataclass` ¶

Loss curve of a short training run (losses[0] is before the first step).

Runs from :func:run_training also hold per-step telemetry (records), the trained setup and the optimizer coverage problems found before training, and offer assertions: :meth:assert_healthy, :meth:assert_finite, :meth:assert_no_skipped_updates, :meth:assert_grad_norm_below, :meth:assert_lr and :meth:assert_loss_decreases. Each returns the run, so they chain.

table ¶

table(highlight: Sequence[int] = (), limit: int = 12) -> str

Per-step telemetry; long runs show the first and last steps plus highlight.

assert_finite ¶

assert_finite() -> TrainingRun

Losses, the gradients applied by every optimizer step and the final parameters are finite.

assert_no_skipped_updates ¶

assert_no_skipped_updates(*, every: int = 1, max_skipped: int = 0) -> TrainingRun

Every every-th step updates the parameters: some optimizer steps, and every optimizer that steps changes them. Up to max_skipped steps may skip the update (e.g. early GradScaler overflows).

assert_grad_norm_below ¶

assert_grad_norm_below(max_norm: float) -> TrainingRun

The total gradient norm stays below max_norm at every optimizer step.

assert_lr ¶

assert_lr(expected: Callable[[int], float] | Sequence[float], *, optimizer: str | None = None, group: int = 0, rtol: float = 1e-06) -> TrainingRun

The learning rate used at each step matches expected: a function of the step index or one value per step. optimizer names it when there are several.

assert_loss_decreases ¶

assert_loss_decreases(min_relative_decrease: float = 0.1) -> TrainingRun

The mean loss of the last quarter of the steps is min_relative_decrease below that of the first quarter (averaging smooths out batch-to-batch noise).

assert_healthy ¶

assert_healthy(*, max_grad_norm: float | None = None, min_relative_decrease: float | None = None, every: int = 1, max_skipped: int = 0) -> TrainingRun

Optimizer coverage, finite values and no skipped updates; plus a gradient-norm bound and a loss decrease when given.

assert_gradients_finite ¶

assert_gradients_finite(model, inputs: Any, loss_fn: LossFn) -> torch.Tensor

Run one backward pass and assert the loss and every gradient are finite.

The model's .grad fields are overwritten. Returns the loss.

assert_all_params_receive_grad ¶

assert_all_params_receive_grad(model, inputs: Any, loss_fn: LossFn, *, ignore: Sequence[str] = (), allow_zero: bool = False) -> None

Assert every trainable parameter gets a gradient from loss_fn.

Catches layers that are defined but never used, outputs detached by mistake (.detach(), .item(), torch.no_grad() in forward) and heads that the loss ignores. A gradient that is exactly zero everywhere also fails unless allow_zero=True. ignore takes fnmatch patterns of parameter names, e.g. ["pooler.*"]. The model's .grad fields are overwritten.

assert_frozen ¶

assert_frozen(model, inputs: Any, loss_fn: LossFn, params: str | Sequence[str], *, optimizer: Callable | None = None, steps: int = 1, check_buffers: bool = True) -> None

Assert the parameters matching params stay frozen through real training steps.

params are fnmatch patterns of parameter names, e.g. "encoder.*". A deep copy of the model is trained in training mode for steps steps with optimizer (a fn(params) -> Optimizer; default AdamW with weight decay, over all parameters like a typical training loop). Fails when a frozen parameter receives a gradient, is changed by the optimizer (e.g. weight decay on a zero gradient), or, with check_buffers, when a buffer under the same names changes: a frozen BatchNorm left in training mode still updates its running statistics. The model you pass is never modified.

assert_gradcheck ¶

assert_gradcheck(fn: Callable, inputs: Any, *, eps: float = 1e-06, atol: float | None = None, rtol: float | None = None) -> None

Compare analytical gradients (autograd) with finite differences.

Use it for custom autograd.Functions and hand-written backward passes. fn is a function or module; it runs in float64 (modules are deep-copied and converted). Gradients are checked with respect to the floating-point input tensors. Keep the inputs small: the cost grows with their number of elements.

assert_loss_decreases ¶

assert_loss_decreases(model, inputs: Any, loss_fn: LossFn, *, steps: int = 20, min_relative_decrease: float = 0.1, lr: float = 0.001, optimizer: Callable | None = None) -> TrainingRun

Assert a few optimizer steps on one batch lower the loss by min_relative_decrease.

A fast smoke test for "the training loop is wired up": wrong sign, frozen parameters, a learning rate of zero or a broken optimizer all fail it. optimizer is fn(params) -> torch.optim.Optimizer (default: Adam with lr). Trains a copy in training mode. Returns the :class:TrainingRun.

assert_overfits_one_batch ¶

assert_overfits_one_batch(model, inputs: Any, loss_fn: LossFn, *, steps: int = 200, max_final_ratio: float = 0.05, target_loss: float | None = None, lr: float = 0.01, optimizer: Callable | None = None) -> TrainingRun

Assert the model can memorize a single batch (the classic first debugging step).

Passes once the loss drops to target_loss or, by default, to max_final_ratio x the initial loss. A model that can't overfit one small batch has a bug (or no capacity): wrong labels alignment, a too-strong regularizer, a disconnected layer... Trains a copy in training mode, stopping early on success. Returns the :class:TrainingRun.

Training loops¶

Training loops: run your real training step for a few steps and check what it does.

A setup is what your training code builds: model, optimizers, schedulers, a GradScaler and any other stateful objects. A step is your training step, step(setup, batch), which does the forward pass, backward pass and optimizer and scheduler steps, and returns the loss (a tensor, a float, or a dict of named losses)::

def make_setup():
    model = Net()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    return TrainingSetup(model, optimizer, schedulers=torch.optim.lr_scheduler.StepLR(optimizer, 10))

def train_step(setup, batch):
    x, y = batch
    loss = F.cross_entropy(setup.model(x), y)
    loss.backward()
    setup.optimizer.step()
    setup.optimizer.zero_grad()
    setup.scheduler.step()
    return loss

run_training(make_setup, batches, train_step, steps=20).assert_healthy()

Everything here builds fresh setups with make_setup and seeds before each run, so runs are reproducible and comparable.

TrainingSetup `dataclass` ¶

Everything a training step uses.

optimizers and schedulers take one object, a list, or a dict of named ones (stored as dicts). state holds other objects with state_dict() and load_state_dict() that a checkpoint must carry, such as an EMA model. step_index is the index of the current step (set by the runner).

optimizer `property` ¶

optimizer: Any

The only optimizer.

scheduler `property` ¶

scheduler: Any

The only scheduler.

TrainingRunner ¶

The training_runner fixture: the functions of this module, bound to the test's torch_device and torch_seed.

run ¶

run(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, **kwargs: Any) -> TrainingRun

:func:run_training.

assert_resume_equivalent ¶

assert_resume_equivalent(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, **kwargs: Any) -> ResumeReport

:func:assert_resume_equivalent.

assert_accumulation_equivalent ¶

assert_accumulation_equivalent(make_setup: Callable[..., Any], data: Any, make_step: Callable[[int], StepFn] | None = None, **kwargs: Any) -> None

:func:assert_accumulation_equivalent.

assert_mixed_precision_training ¶

assert_mixed_precision_training(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, **kwargs: Any) -> TrainingRun

:func:assert_mixed_precision_training.

ResumeReport `dataclass` ¶

Result of :func:assert_resume_equivalent.

run_training ¶

run_training(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, *, loss_fn: Callable[[Any, Any], Tensor] | None = None, steps: int = 20, device: Any = 'cpu', seed: int = 0, track_updates: bool = True) -> TrainingRun

Run steps training steps and record telemetry for each.

make_setup() (or make_setup(device)) returns a :class:TrainingSetup or a (model, optimizer[, scheduler]) tuple; the model is moved to device and put in training mode. data is one batch, a list of batches (cycled), an iterable such as a DataLoader, or a function step_index -> batch. step(setup, batch) is your training step; or pass loss_fn(model, batch) to use :func:standard_step.

Returns a :class:TrainingRun with losses, per-step records (losses, gradient norm, update size and learning rates per optimizer, optimizer step counts, non-finite gradients, GradScaler scale), the trained setup and the optimizer coverage problems, and assertions such as run.assert_healthy(). With track_updates=False parameters aren't copied before each optimizer step (no update_norm), which saves memory on big models.

standard_step ¶

standard_step(loss_fn: Callable[[Any, Any], Tensor], *, accumulation_steps: int = 1, max_grad_norm: float | None = None) -> StepFn

A conventional training step, for when you don't have your own.

loss_fn(model, batch) returns the loss. Gradients accumulate over accumulation_steps steps (the loss is divided accordingly); then every optimizer steps (through setup.scaler when there is one, with gradients clipped to max_grad_norm after unscaling), gradients are reset and every scheduler steps.

assert_optimizer_coverage ¶

assert_optimizer_coverage(model: Any, optimizers: Any, *, ignore: Sequence[str] = (), allow_frozen: bool = True) -> None

Assert every trainable parameter is in exactly one optimizer.

Fails on trainable parameters that no optimizer updates, parameters updated by several optimizers or parameter groups, and optimizers holding tensors that aren't the model's parameters (built before a deepcopy, re-creation or reload of the model). With allow_frozen=False, frozen parameters inside an optimizer fail too. optimizers is one optimizer, a list or a dict; ignore takes fnmatch patterns of parameter names.

assert_resume_equivalent ¶

assert_resume_equivalent(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, *, loss_fn: Callable[[Any, Any], Tensor] | None = None, steps: int = 8, resume_at: int | None = None, save: Callable[[TrainingSetup, Path], Any] | None = None, load: Callable[[TrainingSetup, Path], Any] | None = None, device: Any = 'cpu', seed: int = 0, rtol: float | None = None, atol: float | None = None) -> ResumeReport

Assert that stopping, checkpointing and resuming training changes nothing.

Trains steps steps without interruption, saving a checkpoint with save(setup, path) after resume_at steps (default: half). Then builds a fresh setup, restores the checkpoint with load(setup, path) and trains the remaining steps. Pass your real checkpoint functions; the defaults (:func:save_training_state, :func:load_training_state) save everything.

Fails when the resumed losses or final weights differ from the uninterrupted run, and when the state right after loading differs from the state when saving: model weights, optimizer state and learning rates, scheduler, GradScaler, setup.state. Random number generators that weren't restored are reported as the likely cause when the runs diverge (and ignored when they don't, e.g. without dropout).

training_state ¶

training_state(setup: TrainingSetup) -> dict[str, Any]

Everything needed to resume training exactly: model, optimizers, schedulers, scaler, setup.state and the random number generators.

save_training_state ¶

save_training_state(setup: TrainingSetup, path: str | Path) -> None

Save :func:training_state to path: a reference checkpoint function.

load_training_state ¶

load_training_state(setup: TrainingSetup, path: str | Path) -> None

Restore a checkpoint written by :func:save_training_state into setup.

assert_accumulation_equivalent ¶

assert_accumulation_equivalent(make_setup: Callable[..., Any], data: Any, make_step: Callable[[int], StepFn] | None = None, *, loss_fn: Callable[[Any, Any], Tensor] | None = None, accumulation_steps: int = 4, cycles: int = 2, split: Callable[[Any, int], list[Any]] | None = None, disable_dropout: bool = True, device: Any = 'cpu', seed: int = 0, rtol: float | None = None, atol: float | None = None) -> None

Assert gradient accumulation over micro-batches matches one step on the full batch.

make_step(accumulation_steps) returns your training step configured for that many accumulation steps (it can read setup.step_index); or pass loss_fn(model, batch) to test :func:standard_step. For each of cycles batches, one step with make_step(1) on the full batch is compared with accumulation_steps steps of make_step(accumulation_steps) on its micro-batches (split(batch, n), default :func:split_batch), from identical fresh setups.

Fails on optimizer steps before the last micro-batch or missing after it, and on gradients or weights that differ, saying when gradients are accumulation_steps times too large or too small (loss not divided, or divided twice) and when they carry over from the previous cycle (no zero_grad() after the step). Dropout is disabled in both runs by default, since it draws different masks for a split batch.

split_batch ¶

split_batch(batch: Any, parts: int) -> list[Any]

Split every tensor whose first dimension is the batch size into parts equal micro-batches; other values are passed to each micro-batch unchanged.

assert_mixed_precision_training ¶

assert_mixed_precision_training(make_setup: Callable[..., Any], data: Any, step: StepFn | None = None, *, loss_fn: Callable[[Any, Any], Tensor] | None = None, steps: int = 20, max_skipped_fraction: float = 0.25, device: Any = 'cpu', seed: int = 0) -> TrainingRun

Run a mixed-precision training step and check the GradScaler is used correctly.

Fails when:

optimizer.step() is called directly instead of scaler.step(optimizer);
scaler.step() runs without scaler.scale(loss) or without scaler.update();
gradients are clipped (torch.nn.utils.clip_grad_norm_/clip_grad_value_) before scaler.unscale_(optimizer);
non-finite gradients reach an optimizer step, or the loss itself isn't finite;
more than max_skipped_fraction of the steps are skipped after overflows;
parameters aren't finite after training.

Without a GradScaler (setup.scaler, e.g. bfloat16 autocast), only the last three apply. Returns the :class:TrainingRun.

Performance¶

Performance measurement (latency, peak memory) and baseline tracking.

PerfStore ¶

JSON file with one {latency_ms, peak_memory_mb} entry per name|hardware.

PerfResult `dataclass` ¶

measure_latency ¶

measure_latency(fn: Callable, *args, device: Any = 'cpu', warmup: int = 5, iters: int = 20, **kwargs) -> LatencyStats

Median/p90/min/max wall time of fn(*args, **kwargs) with warmup and device sync.

measure_peak_memory ¶

measure_peak_memory(fn: Callable, *args, device: Any = 'cpu', **kwargs) -> float | None

Peak extra allocated memory in MB during one call. CUDA and XPU; None elsewhere.

check_perf ¶

check_perf(store: PerfStore, key: str, fn: Callable, *args, device: Any = 'cpu', update: bool = False, strict: bool = False, latency_tolerance: float = 0.25, latency_slack_ms: float = 0.05, memory_tolerance: float = 0.1, memory_slack_mb: float = 1.0, warmup: int = 5, iters: int = 20) -> PerfResult

Measure fn(*args) and fail if it regressed against the stored baseline.

A regression is new > baseline * (1 + tolerance) + slack (the slack keeps tiny absolute differences from failing). Faster/leaner is never a failure. A missing baseline skips the test (or fails with strict=True); update=True records one. Peak memory is tracked on CUDA and XPU only.

compare_perf ¶

compare_perf(store: PerfStore, key: str, hardware: str, latency_ms: float, peak_memory_mb: float | None, *, update: bool = False, strict: bool = False, latency_tolerance: float = 0.25, latency_slack_ms: float = 0.05, memory_tolerance: float = 0.1, memory_slack_mb: float = 1.0) -> PerfResult

Compare measured numbers with the baseline for key on hardware (see :func:check_perf), or record them with update=True.

Benchmark suites¶

Benchmark suites: a model's latency, throughput and memory across a matrix of batch sizes, dtypes and execution modes, with per-machine baselines and an end-of-run report.

class TestResNetBenchmark(BenchmarkSuite):
    batch_sizes = (1, 8, 32)
    dtypes = (torch.float32, torch.bfloat16)
    modes = ("eager", "compile")

    def make_model(self):
        return ResNet()

    def make_inputs(self, batch_size, make_tensor):
        return make_tensor(batch_size, 3, 224, 224)

pytest --torch-benchmark                       # run benchmarks (skipped otherwise)
pytest --torch-benchmark --torch-update-perf   # record this machine's baselines
pytest --torch-benchmark --torch-benchmark-json=bench.json

BenchmarkSuite ¶

Base class for a benchmark suite (it is not collected itself).

Override :meth:make_model and :meth:make_inputs, and set the matrix with batch_sizes, dtypes and modes ("eager", "compile"). Every combination becomes one test on every device, measuring median and p90 latency, throughput, peak memory (CUDA, XPU) and, in "compile" mode, compilation time. Latency is measured rounds times (warmup + iters calls each) and the fastest round is kept; a round whose p90 exceeds twice its median is flagged noisy. Results are compared with this machine's baseline (record with --torch-update-perf; tolerances: latency_tolerance, memory_tolerance). Benchmarks only run with --torch-benchmark.

make_inputs ¶

make_inputs(batch_size: int, make_tensor) -> Any

Inputs for one batch of batch_size (tensor, tuple of args or dict of kwargs).

BenchmarkResult `dataclass` ¶

One measured case; collected for the end-of-run report and --torch-benchmark-json.

Compilation and host syncs¶

torch.compile health checks: graph breaks and recompilations.

assert_no_graph_breaks ¶

assert_no_graph_breaks(model, inputs: Any, *, backend: str = 'eager') -> None

Assert the model compiles as ONE graph (fullgraph=True).

Uses the cheap eager backend by default: graph breaks are decided while tracing, so no C++ compiler is needed. On failure the message carries the reason and the offending source line reported by torch.

assert_no_recompiles ¶

assert_no_recompiles(model, inputs: Any, *more_inputs: Any, max_recompiles: int = 0) -> int

Assert repeated calls (and any more_inputs) don't trigger recompilation.

Compiles with a counting backend, runs inputs once, then runs inputs again plus every item of more_inputs and counts how many extra graphs were compiled. A changed tensor shape normally costs ONE recompilation (dynamic shapes kick in), so use max_recompiles=1 when passing different batch sizes. Returns the recompile count.

Detect unexpected host-device synchronization (CUDA).

forbid_host_sync ¶

forbid_host_sync(mode: str = 'error')

Inside the block, synchronizing CUDA ops (.item(), .cpu(), if tensor:) raise (mode="error") or warn ("warn"). Skips the test when CUDA is unavailable.

assert_no_host_sync ¶

assert_no_host_sync(fn: Callable, *args, **kwargs) -> Any

Call fn and fail if it triggers a host-device synchronization.

The original error (with the offending line in its traceback) is chained. Returns fn's result.

Nondeterminism¶

Quantify and diagnose nondeterminism: RNG usage vs. kernel-level variance.

measure_nondeterminism ¶

measure_nondeterminism(model, inputs: Any, *, runs: int = 10, reseed: int | None = None, deterministic: bool = False) -> VarianceResult

Run the model runs times and measure how much the outputs differ from run 1.

reseed reseeds all RNGs before every run (isolates kernel-level variance from RNG usage); deterministic=True runs under torch.use_deterministic_algorithms(True).

diagnose_nondeterminism ¶

diagnose_nondeterminism(model, inputs: Any, *, runs: int = 10, seed: int = 0) -> NondeterminismReport

Find out why outputs vary: RNG usage, kernel-level nondeterminism, or neither.

assert_deterministic ¶

assert_deterministic(model, inputs: Any, *, runs: int = 5, reseed: int | None = None, max_abs: float = 0.0) -> None

Assert the outputs vary by at most max_abs (default: bit-identical) over runs.

Hypothesis strategies¶

Hypothesis strategies for tensors. Needs pip install 'pytest-torch[hypothesis]'.

from hypothesis import given
from pytest_torch.strategies import tensors

@given(x=tensors(shape=(4, 16)))
def test_model_is_finite(x): ...

Note: Hypothesis rejects function-scoped fixtures in @given tests by default; build the model at module scope or suppress HealthCheck.function_scoped_fixture.

shapes ¶

shapes(*, min_dims: int = 1, max_dims: int = 3, min_side: int = 1, max_side: int = 8) -> st.SearchStrategy

Tuples of positive ints (tensor shapes).

tensors ¶

tensors(draw, shape=None, *, dtype: dtype = torch.float32, low: float = -10.0, high: float = 10.0, device: str = 'cpu', edge_cases: bool = True, allow_nan: bool = False, allow_inf: bool = False, noncontiguous: bool = False)

Tensors of dtype with values in [low, high].

shape may be a tuple, a strategy of shapes, or None (random 1-3 dims). With edge_cases=True the data is sometimes all zeros, constant, at the extreme values, or tiny; allow_nan / allow_inf additionally inject one non-finite value. With noncontiguous=True the memory layout is drawn too: contiguous, non-contiguous (like a transposed tensor) or, for 4-D shapes, channels_last. Data comes from a drawn seed, so examples are cheap to generate and replay.

padded_batches ¶

padded_batches(draw, *, vocab_size: int, batch_size: int | SearchStrategy | None = None, max_len: int = 16, min_len: int = 1, pad_id: int = 0, side: str = 'right', ids_key: str = 'input_ids', mask_key: str = 'attention_mask')

{"input_ids", "attention_mask"} batches with ragged, padded sequences.

batch_size is an int or a strategy (default: 1 to 4). Lengths cover the edge cases on purpose: every sequence of length min_len, one full sequence among short ones, all full (no padding), and arbitrary mixes. Real tokens are drawn from [0, vocab_size) excluding pad_id; side is "right" or "left".

Utilities¶

Shared helpers: pytrees, batching, and model device/dtype inspection.

seed_everything ¶

seed_everything(seed: int) -> None

Seed random, numpy (if installed) and torch (CPU, MPS and every CUDA/XPU device).

deterministic_algorithms ¶

deterministic_algorithms(enabled: bool = True)

Temporarily toggle torch.use_deterministic_algorithms; restores the old value.

is_compile_supported ¶

is_compile_supported() -> bool

Whether torch.compile (dynamo) works on this platform and Python version.

Use it to skip compile/export tests: @pytest.mark.skipif(not is_compile_supported(), ...).

Default tolerances: built-in, from pyproject.toml, and per-test overrides.

[tool.pytest_torch.tolerance]
bfloat16 = { rtol = 2e-2, atol = 1e-2 }        # any comparison of bf16 tensors
"cuda.float32" = { rtol = 1e-4 }                # only when testing on CUDA (TF32)
compile_parity = { rtol = 1e-3, atol = 1e-3 }   # default of assert_compile_parity

A key is a dtype name, a check name (see CHECK_DEFAULTS), or either prefixed with a device type (cuda., xpu., mps., cpu.). An entry may set only rtol or only atol; the other comes from the next layer. Resolution, per value:

explicit argument > @pytest.mark.torch_tolerance > pyproject.toml > built-in default

and within a layer the device-specific key beats the generic one. The device is the test's torch_device when the test has one, else the device of the tensors or model compared.

This module must stay torch-free at import: the plugin parses the config at startup.

tolerance_overrides ¶

tolerance_overrides(**entries: Any) -> Iterator[None]

Temporarily override default tolerances, e.g. with tolerance_overrides(bfloat16=(3e-2, 1e-2), compile_parity={"atol": 1e-3}): ....

Keys as in pyproject.toml; use **{"cuda.float32": (1e-4, 1e-5)} for device keys. This is what @pytest.mark.torch_tolerance does for one test, class or module.

Device discovery and per-backend helpers: cpu, cuda, xpu (Intel GPUs) and mps.

torch.cuda and torch.xpu share one API (is_available, device_count, synchronize, memory stats, ...), so GPU code paths look the backend up by device type instead of hard-coding CUDA. Functions an older torch lacks are treated as unsupported.

available_devices ¶

available_devices() -> list[str]

["cpu", ...] plus every accelerator usable in this process.

ModelLoader ¶

Loads models in eval mode and caches them, keyed by every argument.

Cached instances are shared: don't mutate them (use copy.deepcopy first), and note that pretrained weights are downloaded on first use.

hf ¶

hf(name: str, *, auto_class: str = 'AutoModel', dtype: Any = None, device: Any = None, **kwargs) -> Any

transformers.<auto_class>.from_pretrained(name, **kwargs) (path or hub id).

hf_tiny ¶

hf_tiny(name: Any, *, auto_class: str = 'AutoModel', layers: int = 2, hidden_size: int = 64, heads: int = 4, vocab_size: int | None = None, seed: int = 0, dtype: Any = None, device: Any = None, trust_remote_code: bool = False, **overrides) -> Any

A small, randomly initialized model with the architecture of name.

auto_class is an Auto class (default AutoModel) or a model class name such as "LlavaForConditionalGeneration".

Only the config is loaded (name is a hub id, a local path or a PretrainedConfig); :func:shrink_config reduces it to layers layers of width hidden_size with heads attention heads, then overrides set config attributes. Weights are initialized from seed, so the model is the same in every run whatever the test order. model.config is the shrunk config.

timm ¶

timm(name: str, *, pretrained: bool = False, dtype: Any = None, device: Any = None, **kwargs) -> Any

timm.create_model(name, pretrained=..., **kwargs); random init by default.

torchvision ¶

torchvision(name: str, *, weights: Any = None, dtype: Any = None, device: Any = None, **kwargs) -> Any

torchvision.models.get_model(name, weights=...); random init by default.

shrink_config ¶

shrink_config(config: Any, *, layers: int = 2, hidden_size: int = 64, heads: int = 4, vocab_size: int | None = None, **overrides) -> Any

Shrink a Hugging Face PretrainedConfig in place (and return it) for fast tests.

Sets the layer counts, hidden size, attention heads (key/value heads too, keeping grouped-query attention grouped), head dimension, feed-forward sizes and expert counts under every name architectures use for them, truncates per-layer lists (such as layer_types) and recurses into nested configs (text_config, vision_config, encoder, decoder). The vocabulary is kept unless vocab_size is given, so a real tokenizer's ids stay valid. overrides are set last.

BatchFactory ¶

The make_batch fixture: tokens, images and loader.

tokens ¶

tokens(batch_size: int, seq_len: int, *, vocab_size: int = 1000, padding: str | None = 'right', min_length: int | None = None, pad_token_id: int = 0, labels: int | str | None = None, token_type_ids: bool = False) -> dict[str, torch.Tensor]

Token ids and attention mask for sequences of different lengths.

Lengths are drawn between min_length (default half of seq_len) and seq_len; the first row always has the full length and, with more than one row, the last row has min_length, so a batch always mixes padded and unpadded rows. padding is "right", "left" or None (every row full length). Real tokens never equal pad_token_id; padding positions hold it and are 0 in attention_mask. labels adds "labels": an int for that many classes per row, or "lm" for the ids with padding set to -100 (the ignore index).

images ¶

images(batch_size: int, channels: int = 3, height: int = 224, width: int | None = None, *, uint8: bool = False, layout: str = 'nchw', dtype: dtype | None = None, labels: int | None = None) -> Any

A batch of images: normalized floats (about N(0, 1), like ImageNet-normalized pixels) or, with uint8=True, raw pixels in [0, 255]. layout is "nchw" or "nhwc" (raw images as decoded). With labels (a number of classes), returns (images, labels).

loader ¶

loader(data: Any, batch_size: int = 8, *, shuffle: bool = False, **kwargs: Any) -> DataLoader

A DataLoader whose shuffling is seeded by the test's seed (single-process by default). data is a Dataset, a tensor, a tuple of tensors (rows of a TensorDataset) or a dict of equally long tensors (each item a dict).

from_lightning ¶

from_lightning(make_module: Callable[[], Any], *, epoch_length: int | None = None) -> tuple[Callable[[], TrainingSetup], StepFn]

(make_setup, step) for the training-loop checks, from a LightningModule.

make_setup builds the module and calls its configure_optimizers(); step does what Lightning's automatic optimization does for one batch: training_step, backward, optimizer step, zero grad, and the schedulers whose interval is "step" (every frequency steps). There are no epochs in a step-based run: pass epoch_length (steps per epoch) to step "epoch" schedulers too. training_step runs without a Trainer: self.log is ignored and self.global_step reads 0; steps that need self.trainer need the lightning_trainer fixture instead. Manual optimization isn't supported.

make_setup, step = from_lightning(lambda: LitClassifier(lr=1e-3))
training_runner.run(make_setup, batches, step).assert_healthy()

API¶

Generated suites¶

ModelSuite ¶

make_model ¶

make_inputs ¶

loss ¶

frozen_parameters ¶

load_config ¶

config_id classmethod ¶

Assertions¶

assert_tensors_close ¶

assert_outputs_close ¶

assert_finite ¶

assert_batch_invariant ¶

Output specifications¶

TensorSpec dataclass ¶

Probabilities dataclass ¶

assert_output_spec ¶

Checkpoints¶

assert_checkpoint_roundtrip ¶

assert_checkpoint_loads ¶

Stateful and streaming models¶

StreamingEquivalence dataclass ¶

StreamingReport dataclass ¶

assert_streaming_equivalent ¶

Dataset evaluation¶

Evaluation dataclass ¶

PairedDifference dataclass ¶

MetricResult dataclass ¶

evaluate ¶

evaluate_detailed ¶

paired_difference ¶

check_metrics ¶

comparability_problems ¶

model_fingerprint ¶

Failure reproduction¶

record_repro ¶

load_repro ¶

Inference contract¶

ContractReport dataclass ¶

CheckResult dataclass ¶

check_inference_contract ¶

Layer divergence¶

DivergenceReport dataclass ¶

LayerDiff dataclass ¶

localize_divergence ¶

Invariance and layouts¶

assert_invariant ¶

assert_equivariant ¶

assert_permutation_invariant ¶

assert_permutation_equivariant ¶

assert_padding_invariant ¶

assert_mask_invariant ¶

noncontiguous ¶

assert_layout_invariant ¶

Parity¶

assert_compile_parity ¶

assert_torchscript_parity ¶

assert_onnx_parity ¶

assert_export_parity ¶

assert_aoti_parity ¶

assert_autocast_parity ¶

compare_outputs ¶

assert_models_agree ¶

assert_matches_reference ¶

assert_generation_equal ¶

Training¶

StepRecord dataclass ¶

loss property ¶

skipped property ¶

TrainingRun dataclass ¶

table ¶

assert_finite ¶

assert_no_skipped_updates ¶

assert_grad_norm_below ¶

assert_lr ¶

assert_loss_decreases ¶

assert_healthy ¶

assert_gradients_finite ¶

assert_all_params_receive_grad ¶

config_id `classmethod` ¶

TensorSpec `dataclass` ¶

Probabilities `dataclass` ¶

StreamingEquivalence `dataclass` ¶

StreamingReport `dataclass` ¶

Evaluation `dataclass` ¶

PairedDifference `dataclass` ¶

MetricResult `dataclass` ¶

ContractReport `dataclass` ¶

CheckResult `dataclass` ¶

DivergenceReport `dataclass` ¶

LayerDiff `dataclass` ¶

StepRecord `dataclass` ¶

loss `property` ¶

skipped `property` ¶

TrainingRun `dataclass` ¶

TrainingSetup `dataclass` ¶

optimizer `property` ¶

scheduler `property` ¶

ResumeReport `dataclass` ¶

PerfResult `dataclass` ¶

BenchmarkResult `dataclass` ¶