Testing

substrax.testing holds the test infrastructure JAX packages share: fresh-interpreter runs with a chosen JAX configuration, trace, compile and work counts, example runs, and an opt-in pytest plugin. Install it with substrax[testing].

from substrax.runtime import JaxRuntime
from substrax.testing import run_python

result = run_python(
    "import json, jax; print(json.dumps(jax.device_count()))",
    runtime=JaxRuntime(cpu_devices=8),
    timeout=120,
)
assert result.check().last_json() == 8

run_python runs this interpreter on code or a script. The child's environment drops every JAX_* and XLA_* variable the parent inherited, then applies env and the runtime. Unless either of them chooses otherwise, the child runs on the CPU backend without preallocation, so a test never starts an accelerator it did not ask for. check() raises ChildFailedError with the end of the child's standard error, and last_json() parses the last line the child printed. last_json_as(shape) parses it as a value of shape, a TypedDict or a type such as dict[str, list[str]], validated by pydantic in strict JSON mode, so a test reads a structured report with its fields typed and a report of another shape fails naming the field. cuda_is_visible() asks a child on the CUDA backend whether jax sees a GPU.

Counting traces

TraceCounter counts how often a function's Python body runs. Under jax.jit and nnx.jit the body runs once per trace, so the count is the number of traces a test caused:

import jax
import jax.numpy as jnp

from substrax.testing import TraceCounter


def double(x):
    return x * 2.0


counter = TraceCounter()
step = jax.jit(counter.wrap(double))

with counter.expect(new_traces=1):
    step(jnp.ones(3))
with counter.expect(new_traces=0):  # same shape and dtype: no new trace
    step(jnp.zeros(3))

expect raises RetraceError, an AssertionError that names the expected and observed counts, and an error raised inside the block propagates unchanged. Each counter keeps its own count, so there is no global registry to clear. Wrap the Python function, then jit the wrapper. The counter does not listen to jax's jax.monitoring trace event: jax records that event only for top-level traces, and its name is not documented.

Counting compiles

A trace and a compile are different events: a trace builds a jaxpr, a compile builds an XLA executable, and a call that hits the executable cache does neither. substrax.testing.compiles records the executables jax builds inside a block, by name:

import jax
import jax.numpy as jnp

from substrax.testing.compiles import compiled_programs, expect_compiles

step, x = jax.jit(lambda v: v * 2.0), jnp.ones(3)  # jnp.ones compiles too: make inputs first

with compiled_programs() as names:
    step(x)
assert names == ["jit(<lambda>)"]
with expect_compiles(0):  # same shape and dtype: the cached executable
    step(x)

It listens to /jax/core/compile/backend_compile_duration, which jax records around compile_or_get_cached when its in-process executable cache misses, so an executable the persistent compilation cache supplies counts as well. The event name is jax's own constant (jax._src.dispatch.BACKEND_COMPILE_EVENT), not a documented name; the package's tests fail if jax stops recording it. The listener lives for the block only, and it sees every compile in the process, including another thread's. expect_compiles raises CompileCountError, an AssertionError with the expected count and the compiled names. The module imports jax, so substrax.testing does not import it.

Counting work

A test that bounds how much work a code path does by timing it measures the machine as much as the code. The functions the path starts and the calls it makes are the same on every machine. counted_calls counts them with sys.monitoring on the block's own threads: the thread that opened it and every thread started after it opened, so work a block hands to a thread it starts is counted too:

import jax

from substrax.testing import by_package, counted_calls, per_iteration
from substrax.testing.jax_calls import JAX_CALL_WEIGHTS

keys = by_package("mypackage", "grain", "jax", named=["mypackage"])


def count(batches: int) -> dict[str, int]:
    run_pass(batches)  # warm: nothing traces or compiles inside the counted pass
    with counted_calls(keys, weights=JAX_CALL_WEIGHTS) as counts:
        run_pass(batches)
    return counts


per_batch = per_iteration(count, short=100, long=200)
assert per_batch["device_put"] == 1  # one placement per batch
assert per_batch["jit_dispatch"] == 1  # one compiled step per batch
assert per_batch["mypackage:Loader.__next__"] == 1

• by_package adds one to a package's name for each function started in its files, and one to "<package>:<qualified name>" for a package in named; anything else counts as other. Packages are matched by the directories importlib finds for them, and a name it cannot find is refused, so a misspelt package fails instead of counting nothing.
• threading and queue are left out by default (excluded): how often a thread waits on a condition depends on timing, which a count must not.
• Automatic garbage collection is off inside the block and one collection runs at its end, still counted: the collector runs by allocation count, so finalizers would otherwise land in or out of the block depending on what ran before it.
• per_iteration(count, short=, long=) returns (count(long) - count(short)) / (long - short) per key, so work done once per pass (setup, a thread's start and close) cancels. A count that does not grow by a whole number per iteration raises UnevenCountError naming it, instead of being averaged.
• weights count calls made from Python code: each weight sees the callable and the first argument the call passes, and adds what it returns. substrax.testing.jax_calls.JAX_CALL_WEIGHTS counts jit_dispatch (calls of a jax.jit-compiled callable, nnx.jit included), device_put (by identity, so an alias counts) and device_put_leaves (the leaves placed). It imports jax, so substrax.testing does not import it.

Code whose calls depend on object addresses varies between passes, and the counter cannot make it deterministic: an __eq__ of a key hashed by id() runs only when a new object reuses a dead one's address, which the allocator decides. per_iteration names such a key (count its package by name with named to see the function), and a keys function that wraps by_package can leave that code object out, as by_package leaves out threading.

A start is a Python function or generator starting; a generator resuming and a C function are not counted, and a count sees nothing a C function does without calling back into Python. A thread already running Python when the block opens is left out: a library's monitor or heartbeat (tqdm's progress-bar monitor wakes every 10 seconds) would otherwise add its calls to whichever pass it woke in, and per_iteration would find the count uneven. Work a block hands to such a thread, a thread pool created before the block for example, is left out with it; create the pool inside the block to count it. A thread is told apart by the frame its Python execution started in, not by its id: CPython gives the id of a thread that ended to the next one it starts.

counted_calls takes a sys.monitoring tool id for the block and frees it at the end, also when the block raises. It tries the ids CPython leaves unassigned first (3 and 4), so coverage.py's sysmon core (id 1) and cProfile (id 2) keep theirs, and raises ToolIdsInUseError naming the holders when none is free. Python 3.12 and 3.13 run it; on a free-threaded build the callbacks run on several threads at once and a lock keeps the counts exact. cProfile has used sys.monitoring since 3.12 and records the calls of every thread, with no way to leave a thread out.

Checking gradients

A test that only checks that a gradient exists, is finite or is non-zero passes on a wrong one. substrax.testing.gradients compares a module's gradient with finite differences, through jax.test_util.check_grads (forward and reverse mode along one random direction), and returns it:

from substrax.testing.gradients import check_input_gradients, check_parameter_gradients

gradient = check_parameter_gradients(module, lambda m: loss(m(x), y))  # in its nnx.Param leaves
input_gradient = check_input_gradients(module, lambda m, v: loss(m(v), y), x)

The check runs on a float64 copy of the module, split in tree mode: a derivative summed over many output elements is below a float32 finite difference's resolution (check_grads's default tolerance is 1e-5 in float64 against 2e-3 in float32). The caller's module keeps its dtype and values. A gradient that is zero everywhere is refused unless allow_zero=True, since a loss that ignores what is differentiated passes a finite-difference check. A function with a jax.custom_vjp has no forward mode; check it with modes=("rev",).

Running examples

substrax.examples.discover_examples lists a repository's example scripts, and run_example runs one in a fresh interpreter from the repository root:

from pathlib import Path

import pytest

from substrax.examples import discover_examples
from substrax.testing import run_example

ROOT = Path(__file__).resolve().parents[2]
EXAMPLES = discover_examples(ROOT / "examples")


@pytest.mark.parametrize("path", EXAMPLES, ids=lambda path: path.relative_to(ROOT).as_posix())
def test_example_runs(path: Path, output_dir: Path) -> None:
    run = run_example(
        path, repo_root=ROOT, output_dir=output_dir, timeout=600, call_main=True
    )
    assert run.result.check().returncode == 0
    assert run.summary["relative_l2_error"] < 1e-2

• discover_examples skips every file or directory whose name starts with _, such as __init__.py, _common helpers and _templates. include narrows the list further, for example to a numbering rule.
• call_main=True loads the example under a name other than __main__, calls main() and decodes its return value from JSON into run.summary. 0-d numpy and jax values become Python numbers, and a value JSON cannot hold fails the run, naming its key. call_main=False runs the file as a script. In both modes the example's own directory comes first on sys.path, and it sees no command-line arguments, as in a script run.
• The child runs through run_python, with AVITAI_OUTPUT_DIR set to output_dir. An example that resolves its outputs with resolve_output_dir never writes into the repository, and logging or 64-bit settings made at import end with the child.
• A timeout raises ExampleTimeoutError, an AssertionError naming the budget and the end of the example's standard error. A timeout is a failure, not a skip.
• unavailable_reason(run, signatures) returns the first of a repository's messages found in a failed run's standard error, for skipping an example whose dataset is not configured here.

The pytest plugin

Enable it in the top-level conftest.py:

pytest_plugins = ["substrax.testing.pytest_plugin"]

Provides	Behaviour
@pytest.mark.x64	Runs the test inside jax.enable_x64(True), a thread-local setting
Configuration isolation	Snapshots jax.config.values around every test and fails a test that changed a global value, after setting it back
Collection guard	Fails the session when a module left x64 on at import although JAX_ENABLE_X64 did not ask for it
@pytest.mark.devices(count, kind=None)	Skips unless at least count devices, of kind when given, are visible
@pytest.mark.accelerator(kind=None)	Skips unless jax's default backend is an accelerator, of kind when given
output_dir fixture	A fresh outputs directory under tmp_path, with AVITAI_OUTPUT_DIR pointing at it for the test

The isolation follows jax's own test harness, whose JaxTestCase fails a test that changes a global configuration value. A test whose subject changes the global configuration sets it back with restored_jax_config(), which reports the names it restored:

import jax

from substrax.runtime import JaxRuntime, apply_runtime
from substrax.testing import restored_jax_config


def test_apply_runtime_enables_x64() -> None:
    with restored_jax_config() as changed:
        apply_runtime(JaxRuntime(enable_x64=True))
        assert jax.config.jax_enable_x64

    assert changed == ["jax_enable_x64"]

A malformed devices or accelerator marker is a usage error at collection. No entry point loads the plugin, so installing substrax changes no other project's tests, and importing it imports no jax.

Source scans

substrax.testing.source_scans holds the checks a repository's contract tests run over its own source without importing it. Two import-time side effects change the process for everything imported after the module, test collection included: configuring the root logger and writing the environment.

from pathlib import Path

from substrax.testing.source_scans import (
    configures_logging,
    import_time_lines,
    python_files,
    writes_environment,
)

ROOT = Path(__file__).resolve().parents[1]


def test_no_module_changes_the_process_at_import() -> None:
    modules = python_files(ROOT, ("src", "scripts", "tests"), excluded_parts=("example_data",))
    found = [
        f"{path.relative_to(ROOT)}:{line}"
        for path in modules
        for predicate in (configures_logging, writes_environment)
        for line in import_time_lines(path, predicate)
    ]
    assert found == []

import_time_lines skips function, class and lambda bodies and the if __name__ == "__main__": block, whose code runs only when called or run as a script; != and chained comparisons are not guards. writes_environment matches assignment, deletion and |= on os.environ, its mutating methods, and os.putenv/os.unsetenv.

The mesh scan finds documented jax.make_mesh calls that do not name axis_types. From jax 0.11 jax.make_mesh builds Explicit axes by default, and the backward pass of a batch-sharded step raises under them, so every mesh a reader copies names its axis types:

from substrax.testing.source_scans import documented_texts, unguarded_make_mesh_calls


def test_documented_meshes_name_their_axis_types() -> None:
    texts = documented_texts(ROOT, package="datarax", script_directories=("examples",))
    assert {name: calls for name, text in texts.items() if (calls := unguarded_make_mesh_calls(text))} == {}

documented_texts reads every docstring under src/<package>, the README (or the pages given), every Markdown page under docs and every Python file under the script directories. A location that does not exist raises FileNotFoundError, so a mistyped directory fails the check instead of shrinking what it reads.

substrax.testing

Test helpers for JAX packages: fresh-interpreter runs and an opt-in pytest plugin.

run_python runs code or a script in a child interpreter whose JAX settings the test chooses, with everything the parent's shell exported about jax removed, and returns a ChildResult. cuda_is_visible probes the CUDA backend in such a child. restored_jax_config sets back the global jax configuration a block changed. TraceCounter counts the traces a jitted function takes and raises RetraceError when a block caused an unexpected number. run_example runs one example per child interpreter with its outputs redirected (substrax.examples.discover_examples lists them). The plugin, substrax.testing.pytest_plugin, is enabled from a conftest.py: it adds the x64, devices and accelerator markers and fails a test that changes jax's global configuration, as jax's own test harness does. counted_calls counts the Python functions a block's own threads start, grouped by by_package, and per_iteration takes the exact work per iteration from two pass lengths, for tests that bound work by count instead of time; substrax.testing.jax_calls adds jax's dispatches and placements to those counts. substrax.testing.compiles records the XLA programs compiled inside a block and fails a block that compiled an unexpected number, and substrax.testing.gradients checks a module's gradients against finite differences in float64 (both import jax, so the package does not import them). substrax.testing.source_scans holds the contract checks a repository runs over its own source: import-time logging and environment writes, and documented meshes without axis types. Importing this package imports no jax.

CountingFailedError()

Bases: RuntimeError

A key or weight function raised inside the block, so the counts are incomplete.

The function's exception propagates from the call it was counting; this error, chained to it, ends a block whose code caught that exception.

Name the incomplete counts in the message.

ToolIdsInUseError(owners)

Bases: RuntimeError

Every sys.monitoring tool id is held by another tool.

Attributes:

Name	Type	Description
owners	dict[int, str]	The name of the tool holding each id.

Record and name the holder of each id.

Parameters:

Name	Type	Description	Default
owners	dict[int, str]	The name of the tool holding each id.	required

UnevenCountError(deltas, iterations)

Bases: AssertionError

A count did not grow by a whole number per iteration between two passes.

Attributes:

Name	Type	Description
deltas	dict[str, int]	Each such key's count in the long pass minus its count in the short pass.
iterations	int	How many more iterations the long pass ran.

Record the uneven deltas and name them in the message.

Parameters:

Name	Type	Description	Default
deltas	dict[str, int]	Each uneven key's count in the long pass minus the short pass's.	required
iterations	int	How many more iterations the long pass ran.	required

ChildFailedError(result)

Bases: AssertionError

A child interpreter exited with a non-zero code.

Attributes:

Name	Type	Description
result	ChildResult	The failed run.

Describe the failed run by its exit code and the end of its standard error.

Parameters:

Name	Type	Description	Default
result	ChildResult	The failed run.	required

ChildResult(*, argv, returncode, stdout, stderr, seconds) dataclass

What a child interpreter did.

Attributes:

Name	Type	Description
argv	tuple[str, ...]	The command line the child ran.
returncode	int	Its exit code.
stdout	str	Everything it wrote to standard output.
stderr	str	Everything it wrote to standard error.
seconds	float	Wall-clock seconds from start to exit.

check()

Return this result when the child exited with code 0.

Returns:

Type	Description
ChildResult	This result.

Raises:

Type	Description
ChildFailedError	If the exit code is not 0; the message ends with the last lines of standard error.

last_json()

Parse the last non-empty line of standard output as JSON.

Returns:

Type	Description
JsonValue	The decoded value; a caller narrows it to the shape it expects.

Raises:

Type	Description
ValueError	If the child wrote no output, or its last line is not JSON.

last_json_as(shape)

Parse the last non-empty line of standard output as JSON of shape.

The line is validated by pydantic in strict JSON mode, so a value of another type is refused rather than converted (an integer is accepted where shape has a float, as JSON does not tell them apart).

Parameters:

Name	Type	Description	Default
shape	type[T]	The type the value must have, such as a TypedDict or list[int].	required

Returns:

Type	Description
T	The decoded value, of shape.

Raises:

Type	Description
ValueError	If the child wrote no output.
ValidationError	If the line is not JSON, or not a value of shape; the message names the offending field.

ExampleRun(*, path, result, summary) dataclass

One run of an example.

Attributes:

Name	Type	Description
path	Path	The example that ran.
result	ChildResult	The child interpreter's exit code, output and duration.
summary	JsonValue	What main() returned, decoded from JSON, for a run that called it and exited 0; None otherwise.

ExampleTimeoutError(path, timeout, stderr)

Bases: AssertionError

An example ran past its budget and was killed.

Attributes:

Name	Type	Description
path	Path	The example.
timeout	float	Its budget in seconds.
stderr_tail	str	The last lines it wrote to standard error before it was killed.

Name the example, its budget and the end of its standard error.

Parameters:

Name	Type	Description	Default
path	Path	The example.	required
timeout	float	Its budget in seconds.	required
stderr	str	Everything it wrote to standard error before it was killed.	required

RetraceError(expected, observed)

Bases: AssertionError

A block caused a different number of traces than the test expected.

Attributes:

Name	Type	Description
expected	int	The number of new traces the test expected.
observed	int	The number of new traces the block caused.

Record both counts and name them in the message.

Parameters:

Name	Type	Description	Default
expected	int	The number of new traces the test expected.	required
observed	int	The number of new traces the block caused.	required

TraceCounter()

Count how often wrapped functions run their Python body.

jax.jit runs a function's body once per trace and reuses the cached trace for later calls with the same static signature, and nnx.jit calls the user function once inside the function jax traces. The count of body executions is therefore the count of traces. An eager call also runs the body and also counts.

Each counter keeps its own count, so tests share no state to clear.

Start at zero.

count property

The number of body executions of every function this counter wrapped.

wrap(function)

Return function counting each run of its body; jit the result, not the original.

The wrapper keeps the function's name, docstring and signature, so static_argnames and other signature-based arguments of a transform still resolve.

Parameters:

Name	Type	Description	Default
function	Callable[P, R]	The Python function to count.	required

Returns:

Type	Description
Callable[P, R]	A function that adds one to the count, then calls function. The count rises
Callable[P, R]	before the call, so a trace that raises still counts.

expect(*, new_traces)

Require exactly new_traces body executions inside the block.

An exception raised inside the block propagates unchanged, without a count check.

Parameters:

Name	Type	Description	Default
new_traces	int	How many new traces the block must cause.	required

Yields:

Name	Type	Description
None	None	Control to the block.

Raises:

Type	Description
ValueError	If new_traces is negative.
RetraceError	If the block caused a different number of traces.

by_package(*packages, named=(), excluded=DEFAULT_EXCLUDED, other='other')

Group function starts by the package whose files hold the code.

A start in a file under one of packages adds one to that package's name; one in a package of named also adds one to "<package>:<qualified name>". A start in a file of an excluded module or package adds nothing, and any other start adds one to other. Packages are matched by the directories importlib finds for them, the longest first, so a subpackage named alongside its parent takes its own files.

Parameters:

Name	Type	Description	Default
*packages	str	Importable names of the packages to count apart.	()
named	Collection[str]	Packages among packages whose starts are also counted by qualified name.	()
excluded	Collection[str]	Modules or packages whose starts are not counted.	DEFAULT_EXCLUDED
other	str	The key of a start in no listed package.	'other'

Returns:

Type	Description
CodeKeys	The keys of each code object, for :func:counted_calls.

Raises:

Type	Description
ValueError	If a package or excluded module cannot be found, or a named package is not one of packages.

counted_calls(keys, *, weights=None)

Count the Python functions started, and weighted calls made, inside the block.

Only the block's own threads are counted: the thread that opened it and every thread started after it opened, so work the block hands to a thread it starts is counted. A thread already running Python when the block opens (a library's monitor or heartbeat, a pool an earlier test left) is not the block's work and is left out, and so is work the block hands to such a thread, for example to a thread pool created before the block.

keys decides which counts each function start adds one to (it runs once per code object). Each of weights sees every call made from Python code, with the callable and the first argument it passes, and adds what it returns to its own key. Events do not fire while a callback runs, so the weights' own calls are not counted.

Automatic garbage collection is off inside the block; one collection runs at its end, still counted, and the collector's previous state is restored after it.

Parameters:

Name	Type	Description	Default
keys	CodeKeys	The keys of each code object's starts, usually from :func:by_package.	required
weights	Mapping[str, CallWeight] | None	Per key, how much each call adds.	None

Yields:

Type	Description
Generator[Counter[str]]	Counter[str]: The counts, which stop growing when the block ends, also when it raises.

Raises:

Type	Description
ToolIdsInUseError	If every sys.monitoring tool id is held by another tool.
CountingFailedError	If keys or a weight raised inside the block and the block's code caught the exception. Once one has raised, nothing more is counted.

per_iteration(count, *, short, long)

Return the work one iteration adds: the slope between a short and a long pass.

count(n) runs a pass of n iterations and returns its counts. Work done once per pass (setup, a thread's start and close) is the same in both passes and cancels. Each count must grow by a whole number per iteration, so a count that varies between passes fails instead of averaging.

Parameters:

Name	Type	Description	Default
count	Callable[[int], Mapping[str, int]]	Runs a pass of the given number of iterations and returns its counts.	required
short	int	Iterations of the first pass, at least 0.	required
long	int	Iterations of the second pass, more than short.	required

Returns:

Type	Description
dict[str, int]	Each key's count added per iteration; keys that add nothing are left out.

Raises:

Type	Description
ValueError	If short is negative or long is not more than short.
UnevenCountError	If a count does not grow by a whole number per iteration.

cuda_is_visible(*, timeout=120.0)

Whether a fresh interpreter can start jax's CUDA backend and sees at least one GPU.

The probe runs in a child without preallocation, so it claims no accelerator memory.

Parameters:

Name	Type	Description	Default
timeout	float	Seconds before the probe is abandoned.	120.0

Returns:

Type	Description
bool	True when the child found a GPU.

run_python(program, *args, timeout, runtime=None, env=None, cwd=None)

Run Python code or a script in a fresh interpreter and return what it did.

jax fixes its backends and device set when they start, so tests of device topology or process configuration need their own interpreter. The child runs this interpreter (sys.executable). Its environment is this process's without any JAX_* or XLA_* variable, so settings exported in a developer's shell cannot change what a test measures, then env, then runtime as :func:~substrax.runtime.runtime_environment renders it. Unless runtime or env chooses them, the child runs on the CPU backend without preallocation, so a test never starts an accelerator it did not ask for.

Parameters:

Name	Type	Description	Default
program	str | Path	Code to run with -c, or the path of a script.	required
*args	str	Arguments passed to the program.	()
timeout	float	Seconds before the child is killed; every run states its budget.	required
runtime	JaxRuntime | None	JAX settings of the child.	None
env	Mapping[str, str] | None	Variables set in the child before runtime is applied.	None
cwd	Path | None	The child's working directory; this process's when None.	None

Returns:

Type	Description
ChildResult	The child's exit code, output and duration. A non-zero exit is returned, not raised.

Raises:

Type	Description
XlaFlagConflictError	If a flag in runtime conflicts with XLA_FLAGS in env; the child is not started.
TimeoutExpired	If the child runs longer than timeout.

run_example(path, *, repo_root, output_dir, timeout, call_main, runtime=None, env=None)

Run one example in a fresh interpreter from the repository root, with outputs redirected.

The child starts through :func:~substrax.testing.run_python, so the JAX settings are the runtime and env given here, not the parent shell's. AVITAI_OUTPUT_DIR is set to output_dir, so an example resolving its outputs with :func:~substrax.artifacts.resolve_output_dir never writes into the repository. Whatever the example configures at import, such as logging or 64-bit types, ends with the child.

Parameters:

Name	Type	Description	Default
path	Path	The example.	required
repo_root	Path	The child's working directory, which relative paths in examples assume.	required
output_dir	Path	Where the example's outputs go.	required
timeout	float	Seconds before the child is killed; every run states its budget.	required
call_main	bool	True loads the example under a name other than __main__, calls its main() without arguments and decodes the return value; False runs the file as a script.	required
runtime	JaxRuntime | None	JAX settings of the child.	None
env	Mapping[str, str] | None	Variables set in the child before runtime is applied.	None

Returns:

Type	Description
ExampleRun	The run. A non-zero exit is returned, not raised; run.result.check() raises it.

Raises:

Type	Description
ValueError	If env sets AVITAI_OUTPUT_DIR to a directory other than output_dir.
ExampleTimeoutError	If the example runs longer than timeout.
XlaFlagConflictError	If a flag in runtime conflicts with XLA_FLAGS in env; the child is not started.

unavailable_reason(run, signatures)

Return the first signature found in a failed run's standard error, ignoring case.

A repository lists the messages that mean an example cannot run in this environment, such as a dataset that is not configured, and skips the example when one matches. A run that exited 0 matches nothing, and standard output is not searched.

Parameters:

Name	Type	Description	Default
run	ExampleRun	The finished run.	required
signatures	Sequence[str]	Messages that mean the example cannot run here.	required

Returns:

Type	Description
str | None	The first signature found, as given, or None.

restored_jax_config()

Set back every jax configuration value a block changed, and report which ones.

jax's own test base class snapshots jax.config.values around each test and fails a test that changed a global value (assert_global_configs_unchanged in jax/_src/test_util.py). This context manager takes the same snapshot and, on exit, sets every value that differs back through jax.config.update, whether the block finished or raised.

A value read inside a thread-local context manager such as jax.enable_x64(True) is that context's, so a context entered and left inside the block changes nothing.

Yields:

Type	Description
list[str]	list[str]: Empty while the block runs; on exit it holds the sorted names of the values the
list[str]	block changed, all of them restored.

Raises:

Type	Description
RuntimeError	If jax refuses to set a value back, as it refuses jax_num_cpu_devices once its backends have started.

substrax.testing.compiles

The XLA executables jax builds inside a block, for tests that a repeated call compiles nothing.

jax records the event /jax/core/compile/backend_compile_duration around compile_or_get_cached, which runs only when its in-process executable cache misses, and passes the program's name, jit(<function>), as fun_name (jax/_src/interpreters/pxla.py and jax/_src/dispatch.py, jax 0.11.1). Eager operations are compiled programs too. A block therefore sees every executable built for it, whether XLA compiled it or the persistent compilation cache supplied it. A trace is a different event: :class:~substrax.testing.TraceCounter counts those.

The listener is registered with jax.monitoring for the block and removed after it. Listeners are process-wide, so a compile on another thread during the block is seen too.

BACKEND_COMPILE_EVENT = '/jax/core/compile/backend_compile_duration' module-attribute

The jax.monitoring duration event jax records for each executable it builds.

CompileCountError(expected, names)

Bases: AssertionError

A block built a different number of executables than the test expected.

Attributes:

Name	Type	Description
expected	int	The number of compiles the test expected.
names	tuple[str, ...]	The names of the functions compiled in the block, in order.

Record the expectation and what was compiled, and name both in the message.

Parameters:

Name	Type	Description	Default
expected	int	The number of compiles the test expected.	required
names	tuple[str, ...]	The names of the functions compiled in the block.	required

compiled_programs()

Collect the names of the functions jax compiles inside the block.

Yields:

Type	Description
Generator[list[str]]	list[str]: The name of each program compiled in the block, in order; the list stops
Generator[list[str]]	growing when the block ends, also when the block raises.

expect_compiles(count)

Require exactly count compiles inside the block.

An exception raised inside the block propagates unchanged, without a count check.

Parameters:

Name	Type	Description	Default
count	int	How many executables the block must build.	required

Yields:

Name	Type	Description
None	Generator[None]	Control to the block.

Raises:

Type	Description
ValueError	If count is negative.
CompileCountError	If the block built a different number of executables.

substrax.testing.jax_calls

Weights for :func:~substrax.testing.counted_calls that count jax's dispatches and placements.

Pass :data:JAX_CALL_WEIGHTS as weights to count, in the block:

• jit_dispatch: calls of a jax.jit-compiled callable from Python code, whether the call hits the executable cache or traces. nnx.jit calls one, and so do the jax.numpy functions jax jits itself. Compiles are a different event: :mod:substrax.testing.compiles counts those.
• device_put: calls of jax.device_put, by identity, so a call through an alias counts.
• device_put_leaves: the leaves of the first argument each jax.device_put call passes, the value placed when it comes first, by position or as x=. A call that passes another keyword before x counts that keyword's leaves.

A call counts only where Python code makes it: a dispatch that jax's C++ code makes, or a placement inside a traced function, fires no call event. The module imports jax, so :mod:substrax.testing does not import it.

JAX_CALL_WEIGHTS = MappingProxyType({'jit_dispatch': jit_dispatches, 'device_put': device_puts, 'device_put_leaves': leaves_placed}) module-attribute

The three jax counters by key, for counted_calls(..., weights=JAX_CALL_WEIGHTS).

jit_dispatches(callee, first)

Count one for each call of a jax.jit-compiled callable.

Parameters:

Name	Type	Description	Default
callee	object	The called object.	required
first	object	The first argument the call passes (unused).	required

Returns:

Type	Description
int	1 for a compiled callable, else 0.

device_puts(callee, first)

Count one for each call of jax.device_put.

Parameters:

Name	Type	Description	Default
callee	object	The called object.	required
first	object	The first argument the call passes (unused).	required

Returns:

Type	Description
int	1 for jax.device_put, else 0.

leaves_placed(callee, first)

Count the leaves of the value each jax.device_put call places.

Parameters:

Name	Type	Description	Default
callee	object	The called object.	required
first	object	The first argument the call passes, the value placed.	required

Returns:

Type	Description
int	The number of leaves of first for a jax.device_put call that passes an argument,
int	else 0.

substrax.testing.gradients

Gradient checks that assert values: a module's autodiff gradient against finite differences.

A test that only checks a gradient exists, is finite or is non-zero passes on a wrong gradient. These checks compare the gradient of a loss, in a module's nnx.Param leaves or in its input, with finite differences through jax.test_util.check_grads (forward and reverse mode along one random direction; a function with a jax.custom_vjp has no forward mode, so it is checked with modes=("rev",)). They run on a float64 copy of the module: a derivative summed over many output elements is below a float32 finite difference's resolution (check_grads's default tolerance is 1e-5 in float64 against 2e-3 in float32). The copy is split in tree mode, which a module holding an nnx.RngKey needs, and the caller's module keeps its own dtype and values.

The module imports jax and flax, so substrax.testing does not import it.

check_parameter_gradients(module, loss, *, modes=('fwd', 'rev'), allow_zero=False)

Check the gradient of loss in module's nnx.Param leaves, and return it.

Parameters:

Name	Type	Description	Default
module	M	The module; it is not changed.	required
loss	Callable[[M], Array]	A scalar function of the module.	required
modes	Sequence[str]	The autodiff modes to check, "fwd" and/or "rev".	('fwd', 'rev')
allow_zero	bool	Accept a gradient that is zero everywhere; refused otherwise, since a loss that does not depend on the parameters passes a finite-difference check.	False

Returns:

Type	Description
NnxState	The gradient, in float64, with the structure of the module's nnx.Param state.

Raises:

Type	Description
AssertionError	If the gradient differs from finite differences, or is zero everywhere and allow_zero is False.

check_input_gradients(module, loss, value, *, modes=('fwd', 'rev'), allow_zero=False)

Check the gradient of loss in its input value, and return it.

Parameters:

Name	Type	Description	Default
module	M	The module; it is not changed.	required
loss	Callable[[M, Array], Array]	A scalar function of the module and the input.	required
value	ArrayLike	The input the gradient is taken at, cast to float64.	required
modes	Sequence[str]	The autodiff modes to check, "fwd" and/or "rev".	('fwd', 'rev')
allow_zero	bool	Accept a gradient that is zero everywhere; refused otherwise.	False

Returns:

Type	Description
Array	The gradient in float64, shaped like value.

Raises:

Type	Description
AssertionError	If the gradient differs from finite differences, or is zero everywhere and allow_zero is False.

substrax.testing.source_scans

Source scans a repository's contract tests run over its own files.

Two families of check read source rather than run it:

• Import-time side effects. :func:import_time_lines walks a module and returns the lines of the nodes a predicate matches, among those that run when the module is imported: function, class and lambda bodies and the if __name__ == "__main__": block are skipped. :func:writes_environment matches a change to os.environ (assignment, deletion, |=, a mutating method, os.putenv/os.unsetenv); :func:configures_logging matches a logging.basicConfig call. A module that does either at import changes the process for everything imported after it, test collection included.
• Documented meshes. :func:unguarded_make_mesh_calls returns each make_mesh call whose arguments do not name axis_types: from jax 0.11 jax.make_mesh builds AxisType.Explicit axes by default, and the backward pass of a batch-sharded step raises under them. :func:documented_texts collects what a reader copies from: every docstring in the package, the named pages, every Markdown page under the page directories and every Python file under the script directories.

A corpus location that does not exist raises :class:FileNotFoundError, so a mistyped directory fails the check instead of shrinking what it reads. Importing this module imports no jax.

python_files(root, directories, *, pattern='*.py', excluded_parts=())

Every file under root / directory matching pattern, in directory then path order.

Parameters:

Name	Type	Description	Default
root	Path	The repository root.	required
directories	Sequence[str]	Directories relative to root to search recursively.	required
pattern	str	The file-name glob.	'*.py'
excluded_parts	Collection[str]	Directory names whose files are skipped, such as "example_data".	()

Returns:

Type	Description
list[Path]	The matching files.

Raises:

Type	Description
FileNotFoundError	If a directory does not exist.

writes_environment(node)

Whether node changes os.environ or calls os.putenv/os.unsetenv.

configures_logging(node)

Whether node calls logging.basicConfig, by module attribute or imported name.

is_main_guard(node)

Whether node is if __name__ == "__main__": (either operand order).

Its body runs only when the module is run as a script. != and chained comparisons are not guards: their bodies run at import.

import_time_lines(path, matches)

Line numbers of the nodes matches accepts that run when path is imported.

Parameters:

Name	Type	Description	Default
path	Path	A Python source file.	required
matches	Callable[[AST], bool]	A predicate over one AST node, such as :func:writes_environment.	required

Returns:

Type	Description
list[int]	The line numbers in source order.

unguarded_make_mesh_calls(text)

Each make_mesh call in text whose arguments do not name axis_types.

A call whose parentheses never close is returned from its name to the end of the text.

documented_texts(root, *, package, pages=('README.md',), page_directories=('docs',), script_directories=())

Map every text a reader copies code from to its content, keyed by a repository path.

Parameters:

Name	Type	Description	Default
root	Path	The repository root.	required
package	str	The package under root / "src" whose docstrings are read.	required
pages	Sequence[str]	Individual pages relative to root.	('README.md',)
page_directories	Sequence[str]	Directories whose Markdown pages are read.	('docs',)
script_directories	Sequence[str]	Directories whose Python files are read whole, such as "examples".	()

Returns:

Type	Description
dict[str, str]	Docstrings keyed "<path>:<line>" (line 0 for a module docstring), and pages and
dict[str, str]	scripts keyed by their path, all relative to root.

Raises:

Type	Description
FileNotFoundError	If the package, a page or a directory does not exist.

substrax.examples

A repository's example scripts: every *.py under its examples directory that is not private.

Importing this module imports no jax.

discover_examples(root, *, include=None)

List the example scripts under root, sorted.

A file or directory whose name starts with _ is private, as _common helpers, _templates and __init__.py are, and so is everything under it. Only the parts below root count, so root itself may have such a name.

Parameters:

Name	Type	Description	Default
root	Path	The examples directory.	required
include	Callable[[Path], bool] | None	Keeps a path when it returns True, such as a repository's numbering rule.	None

Returns:

Type	Description
list[Path]	The *.py files under root that are not private and that include keeps.

substrax.testing.compute

The behaviour every compute backend owes, as pytest tests a backend's own suite inherits.

Subclass :class:BackendContract in a test module and define make_backend; the inherited tests submit small jobs of plain Python scripts, wait for them, and check the states, logs, fetched outputs and cancellation the :class:~substrax.compute.backend.ComputeBackend protocol promises. The scripts import no jax, so a backend that runs them locally finishes each test in seconds.

class TestLocalBackend(BackendContract):
    def make_backend(self, state_dir: Path) -> ComputeBackend:
        return LocalBackend({"python": sys.executable, "state_dir": str(state_dir)})

BackendContract

Tests every backend must pass; subclass and define :meth:make_backend.

make_backend(state_dir)

Build a fresh backend instance keeping whatever it stores under state_dir.

Parameters:

Name	Type	Description	Default
state_dir	Path	A directory private to the test.	required

Returns:

Type	Description
ComputeBackend	The backend.

Raises:

Type	Description
NotImplementedError	Always, until a subclass defines it.

project(tmp_path)

A project holding the contract's scripts.

Parameters:

Name	Type	Description	Default
tmp_path	Path	The test's directory.	required

Returns:

Type	Description
Path	The project root.

state_dir(tmp_path)

The directory the backend under test keeps its state in.

Parameters:

Name	Type	Description	Default
tmp_path	Path	The test's directory.	required

Returns:

Type	Description
Path	The directory.

test_a_job_succeeds_and_its_outputs_come_back(project, state_dir, tmp_path)

A succeeding job ends succeeded and fetches its logs, outputs and manifest.

test_a_failed_task_fails_the_run_and_the_tasks_after_it_still_run(project, state_dir, tmp_path)

A failed task makes the run failed without skipping the rest.

test_following_the_logs_ends_with_the_run_and_names_each_task(project, state_dir)

Followed logs carry each task's lines, prefixed with its name, and stop at the end.

test_a_cancelled_run_ends_cancelled(project, state_dir)

Cancelling a running job ends it as cancelled; cancelling again changes nothing.

test_a_stored_handle_finds_the_run_from_a_new_backend(project, state_dir, tmp_path)

A handle read back from a RunStore is enough for another backend instance.

wait_until_final(backend, handle, *, timeout=_WAIT_SECONDS)

Poll a run until it ends.

Parameters:

Name	Type	Description	Default
backend	ComputeBackend	The backend that runs it.	required
handle	RunHandle	The run.	required
timeout	float	Seconds to wait.	_WAIT_SECONDS

Returns:

Type	Description
RunState	Its final state.

Raises:

Type	Description
TimeoutError	If the run has not ended after timeout seconds.