v1.1.1

1.1.1

A fix for --torch-num-threads=auto under pytest-xdist: it now leaves a thread count you
set with MKL_NUM_THREADS alone, as it already did for OMP_NUM_THREADS.

Fixed

• --torch-num-threads=auto also leaves the thread count alone when MKL_NUM_THREADS is set,
  not only OMP_NUM_THREADS: torch gives MKL_NUM_THREADS precedence, so it's your setting
  too. An explicit count now sets both variables for subprocesses.
• The test of that behavior no longer depends on how torch resolves the two variables on
  the machine running it; it failed on CI runners that set MKL_NUM_THREADS.

Improvements, benchmark and training loops

1.1.0

Testing training code, not just models. You can now run your own training step, checkpoint
code and fine-tuning setup for a few seeded steps, and get each problem named: an optimizer
that misses parameters, state missing from a checkpoint, wrong gradient-accumulation
scaling, GradScaler misuse. This release also adds benchmark suites, tiny versions of
Hugging Face architectures, realistic synthetic batches, test fixtures for the Lightning and
Hugging Face trainers, and markers that replace hand-written skipifs and setup code. It is
backward compatible with 1.0: everything new is additive (see Changed for two
behavior changes).

Training loops

Run your own training step and checkpoint code for a few seeded steps from fresh setups
(TrainingSetup), with the training_runner fixture or the functions directly. See
examples/test_training_loops.py and the new Training loops guide.

• run_training records per-step telemetry (StepRecord: losses, gradient norm, update
  size, learning rates, optimizer steps, non-finite gradients, loss scale) without changes
  to the training code. The returned TrainingRun adds assert_healthy, assert_finite,
  assert_no_skipped_updates, assert_grad_norm_below, assert_lr,
  assert_loss_decreases and table().
• assert_optimizer_coverage: trainable parameters in no optimizer or in several, and
  optimizers built for another copy of the model.
• assert_resume_equivalent: training resumed from a checkpoint matches uninterrupted
  training. The state after loading is compared with the saved state to name what's
  missing (optimizer, scheduler, GradScaler, RNG). save_training_state,
  load_training_state and training_state are reference implementations.
• assert_accumulation_equivalent: gradient accumulation matches one full-batch step,
  diagnosing an unscaled loss, gradients not reset and optimizer steps at the wrong
  micro-batch.
• assert_mixed_precision_training: GradScaler misuse (bypassed scaler.step, missing
  scaler.update, clipping before unscale_), overflows and non-finite values.
• standard_step(loss_fn, accumulation_steps=, max_grad_norm=) for code without its own
  step, and split_batch.
• Tolerance checks training_resume and grad_accumulation.
• TrainingRun gained records, setup, coverage_problems and the assertions above;
  assert_loss_decreases and assert_overfits_one_batch return it as before.

Lightning and the Hugging Face Trainer

• lightning_trainer and hf_trainer fixtures: a Lightning Trainer or a Hugging Face
  Trainer set up for tests: the test's device and seed, 5 steps, no loggers, reports or
  checkpoints, files in tmp_path, any argument overridable. hf_trainer logs the loss
  every step, accepts make_batch.tokens(...) as a dataset and builds your Trainer
  subclass with trainer_class=. See examples/test_hf_trainer.py.
• from_lightning(make_module): the make_setup and step of a LightningModule with
  automatic optimization, for training_runner and the training-loop checks.
• New lightning extra.

Benchmark suites

• Subclass BenchmarkSuite, implement make_model() and
  make_inputs(batch_size, make_tensor), and set batch_sizes, dtypes and modes
  ("eager", "compile"). Every combination runs as one test per device and measures:

  • median and p90 latency, and throughput;
  • peak memory (CUDA, XPU);
  • compilation time in "compile" mode.

  Each case is measured rounds times and the fastest round is kept. Cases are compared
  with per-machine baselines (--torch-update-perf, latency_tolerance,
  memory_tolerance). The run ends with a report table. Benchmarks run only with
  --torch-benchmark; --torch-benchmark-json also writes the results
  (BenchmarkResult). Dtypes the device can't run, and torch.compile where it doesn't
  work, are skipped. See examples/test_benchmarks.py and the benchmark in the image
  classification integration example.

• compare_perf: the baseline comparison of check_perf, for numbers you measured
  yourself.

• LatencyStats.p90_ms.

Synthetic data and tiny models

• make_batch fixture: realistic synthetic batches drawn through make_tensor (seeded,
  on torch_device, recorded in reproduction bundles). tokens() gives input_ids and
  attention_mask for rows of mixed lengths (always one full row and one padded row,
  padding left or right, real tokens never equal to the pad id, optional class or
  language-modeling labels). images() gives normalized or raw uint8 images in NCHW or
  NHWC. loader() gives a DataLoader shuffled by the test's seed.
• model_loader.hf_tiny(name): a small, randomly initialized model with the architecture of
  a Hugging Face model, from its config alone, with weights fixed by a seed. shrink_config
  does the shrinking (layers, widths, attention and key/value heads, feed-forward sizes,
  experts, per-layer lists, nested configs). --torch-model=hf-tiny:NAME uses it for
  zero-code testing: the Qwen2.5-7B architecture goes through the zero-code suite in
  seconds on a CPU.

Markers

• torch_requires(...): skip unless devices (cuda, xpu, mps, gpu), versions
  (torch, python specifiers), compile, gpu_memory_gb, cuda_capability and
  packages are available, with the reason at the test's location.
• torch_skip_on(*devices) and torch_xfail_on(*devices): skip or expect a failure for one
  device (and optionally dtype) only. Expected failures are strict and take raises=.
• torch_deterministic, torch_default_dtype(dtype), torch_matmul_precision(...)
  ("highest" disables TF32, convolutions included), torch_no_host_sync,
  torch_detect_anomaly: set for one test, fixtures included, and restored afterwards.
• torch_max_memory(mb=): peak memory budget of the test body (CUDA, XPU).
• torch_memory_leak(iterations=5): runs the body repeatedly and fails when every run
  leaves more tensors or memory behind, describing the new tensors and their holders.
• torch_threads(n): the number of CPU threads torch uses during one test.
• A new Markers reference page documents every marker with examples.

Parallel runs

• --torch-num-threads (default auto): under pytest-xdist, each worker uses its share of
  the CPU cores instead of one thread per core, unless OMP_NUM_THREADS is set. An integer
  sets the count everywhere; none keeps torch's default.

Changed

• Under pytest-xdist, workers now use their share of the CPU cores (--torch-num-threads=auto)
  instead of torch's default of one thread per core each. Parallel runs no longer
  oversubscribe the machine, but CPU-bound timings change; --torch-num-threads=none
  restores the previous behavior. Runs without xdist are unaffected.
• The hf extra also installs accelerate, which the Hugging Face Trainer needs.

Fixed

• check_inference_contract:
  • copies the first output right away, so a model returning a buffer it overwrites in
    place can't pass the determinism check by aliasing;
  • checks every registered parameter and buffer for mutation, non-persistent buffers
    included, and reports state added or removed during inference (plain Python
    attributes stay outside the check);
  • rejects an empty selection (checks=[], or skipping every check) with a
    ValueError before running the model, and doesn't run the model for eval_mode
    alone.

Development

• CI: an integrations job runs the tests with transformers, timm, torchvision and
  Lightning installed; until now, the tests that need them were skipped in CI.

v1.0.0

What's Changed

• First better release by @AKJ7 in #3
• v1.0.0 by @AKJ7 in #4

Full Changelog: https://github.com/AKJ7/pytest-torch/commits/v1.0.0

1.0.0

First stable release (the earlier 0.1.0.dev0 on PyPI was a preview). From here on,
pytest-torch follows semantic versioning: see
Stability below. pytest-torch is a pytest plugin for testing PyTorch models. It
covers inference correctness, invariance, parity with deployment targets, training sanity,
performance and determinism, and its failure messages say where a mismatch is and how large
it is. It requires Python 3.9+, pytest 7+ and PyTorch 2.0+. The plugin registers itself
when installed, and importing it never imports torch.

Test any model with no code

pytest --torch-model=hf:google-bert/bert-base-uncased
pytest --torch-model=timm:resnet50 --torch-model=mypackage.models:Classifier

• Model specs: hf:NAME[@AutoClass], timm:NAME, torchvision:NAME,
  package.module:Factory or path/to/file.py:Factory. Weights are random by default
  (Hugging Face models only download their config); --torch-model-pretrained loads
  pretrained weights.
• Inferred inputs: padded input_ids/attention_mask for text models, images of the
  configured size for vision models, the model's example_inputs(), or the first layer's
  shape. --torch-model-inputs describes them when inference can't.
• Checks: the inference contract, padding/mask and layout invariance, a checkpoint
  round-trip and finite gradients by default; export and compile on request with
  --torch-model-checks. A model that can't be built fails once, in test_builds, with
  the fix.

A whole test suite from one class: ModelSuite

Subclass ModelSuite, describe how to build the model and its inputs, and pick the checks.
pytest collects one test per check, on every available device, with seeded weights and
inputs.

from pytest_torch import ModelSuite

class TestClassifier(ModelSuite):
    checks = ("contract", "compile", "export", "training")

    def make_model(self):
        return Classifier()

    def make_inputs(self, make_tensor):
        return make_tensor(8, 16)

    def loss(self, output, inputs):          # only needed for "training"
        return F.cross_entropy(output, labels)

• Available checks: contract, compile, export (these three are the default),
  torchscript, onnx, aoti, autocast, golden and training.
• check_options passes arguments to the assertion behind each test, for example
  {"overfits_one_batch": {"steps": 500}}.
• With dynamic_batch = True, the export test checks that one exported program accepts
  other batch sizes.
• Checks you don't list are not generated at all, so there are no skipped placeholder tests.
• Generated tests are ordinary methods: you can override them or add your own.
• Configuration mistakes (an unknown check, an unknown option, training without a
  loss() method) fail at collection with a message saying what to fix.
• One suite per config: with configs = sorted(Path("configs").glob("*.yaml")), every
  generated test runs once per config file (YAML, JSON or Python objects), with
  config_overrides to shrink each config and self.config / the model_config fixture
  to build from it. Ids read test_frozen[transformer_rope_frozen-cpu].
• frozen check: parameters a config freezes must stay frozen through training.

Fixtures and reproducibility

Fixture	What it gives you
torch_device	Runs the test on every available device (CPU, CUDA, Intel XPU, Apple MPS); unavailable ones are skipped
torch_seed	Seeds random, numpy and torch; the seed is printed when the test fails
torch_generator	A seeded CPU generator, so the same seed gives identical values on every device
make_tensor	Reproducible tensors (normal, uniform or integer), drawn on the CPU and moved to the device
inference_mode	Runs the test body under torch.inference_mode()
golden	Snapshot tests of tensors and nested outputs
torch_dtype	Runs the test once per dtype (default float32; --torch-dtype or @pytest.mark.torch_dtypes)
latency, perf_baseline	Latency budgets and performance regression checks
metric_baseline	Dataset metrics compared with a stored baseline and thresholds
model_loader	Session-cached Hugging Face, timm and torchvision models

• Device selection: --torch-device pins one or more devices.
• Seeding: the seed comes from @pytest.mark.torch_seed(n) if set, otherwise from
  --torch-seed-mode=per-test (a stable seed derived from each test's id), otherwise from
  --torch-seed. --torch-autoseed seeds every test, and --torch-deterministic turns on
  torch.use_deterministic_algorithms inside torch_seed.
• Golden files: pytest --torch-update-golden writes them. They store nested outputs
  by path and record the torch version, which is shown when a comparison fails. A missing
  golden file fails the test, so CI can't pass without comparing anything.

Project-wide tolerances

Set default tolerances once in pyproject.toml instead of repeating rtol=/atol= in
every test:

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

• Keys: a dtype, a check (compile_parity, export_parity, device_parity,
  autocast_parity, gradcheck and others), or either prefixed with a device. An entry
  may set only rtol or only atol.
• One test, class or module: @pytest.mark.torch_tolerance(bfloat16=(5e-2, 1e-2)).
  For a block of code, use with tolerance_overrides(...).
• Precedence: an argument passed to the assertion wins, then the marker, then
  pyproject.toml, then the built-in default.
• Mistakes in the config stop pytest at startup with a message saying what is wrong,
  and failure messages always print the tolerance that was used.

Assertions with readable failures

• Comparisons: assert_tensors_close uses default tolerances suited to each dtype.
  Complex tensors compare both real and imaginary parts, and integer and bool tensors
  compare exactly (also above 2^53) unless a tolerance is given. NaNs count as equal unless
  equal_nan=False.
• Nested outputs: assert_outputs_close compares dicts, tuples, dataclasses and Hugging
  Face ModelOutputs, and reports the path of each mismatch, such as output['logits']. It
  also reports structure changes, including empty containers ([] vs {}).
• assert_finite finds NaN or Inf anywhere in a nested output.
• assert_batch_invariant checks that each sample's output doesn't depend on the rest
  of the batch.
• Failures report how many elements differ, where the worst one is, and the absolute and
  relative error.
• Model inputs can be a tensor, a tuple of positional arguments or a dict of keyword
  arguments, and they are moved to the model's device automatically.

Inference contract

check_inference_contract(model, inputs) runs the checks every inference model should
pass, and gives a likely cause for each failure:

• eval_mode: no module was left in training mode.
• finite_outputs: no NaN or Inf in the outputs.
• deterministic: two runs give identical outputs.
• batch_invariance: a sample's output doesn't depend on the rest of the batch.
• device_parity: results match on the other available devices, or only on the other
  devices pinned with --torch-device.
• dtype_parity: results stay close at lower precision. It only runs if you ask for it,
  for example with dtypes=(torch.bfloat16,).
• no_state_mutation: inference didn't change any parameter or buffer.

You can select or skip checks and tune tolerances. With raise_on_failure=False you get a
report back instead of an exception.

Find the first diverging layer

localize_divergence(model_a, model_b, inputs) compares two models layer by layer (for
example CPU against GPU, fp32 against bf16, or a port against the original) and names the
first layer that disagrees. It also tells you whether the error started inside that layer
or came in from the code between layers. Activations are copied when recorded, so in-place
operations later in forward() can't distort the comparison. module_map pairs up layers whose names differ.

Invariance and metamorphic tests

These check properties that must hold even when you don't know the right output:

• assert_padding_invariant: adding padding (left or right) doesn't change the
  outputs for real tokens.
• assert_mask_invariant: changing the tokens under the attention mask doesn't change
  the outputs.
• assert_permutation_invariant: reordering the elements of a set doesn't change the
  output.
• assert_permutation_equivariant: reordering the input reorders the output the same
  way.
• assert_invariant and assert_equivariant: check any relation you define yourself.

Outputs aligned with the sequence are compared only on real tokens.

Parity with deployment targets and references

• torch.compile: assert_compile_parity compares compiled and eager outputs over
  repeated runs.
• torch.export: assert_export_parity compares the exported program with eager
  mode. Pass extra inputs plus dynamic_batch=True, or your own dynamic_shapes, to check
  that one program handles several batch sizes. If export fails, th...