Records

A record is a frozen dataclass written to JSON by substrax.records.dump_record and read back by substrax.records.read_record, which validates every field against its annotation with pydantic's strict JSON mode. substrax.typing holds the aliases records are typed with; it imports nothing beyond the standard library.

from dataclasses import dataclass

from substrax.records import dump_record, read_record, UNKNOWN_FIELDS_REFUSED


@dataclass(frozen=True, slots=True, kw_only=True)
class Measurement:
    __pydantic_config__ = UNKNOWN_FIELDS_REFUSED

    name: str
    value: float
    samples: tuple[float, ...] | None = None


dump_record(Measurement(name="loss", value=0.25, samples=(0.2, 0.3)))
# {'name': 'loss', 'value': 0.25, 'samples': [0.2, 0.3]}

read_record(Measurement, {"name": "loss", "value": "0.25"})
# pydantic.ValidationError: 1 validation error for Measurement
# value
#   Input should be a valid number [type=float_type, input_value='0.25', input_type=str]

Annotation	Reads
str, bool	The JSON value of that type only
int	A JSON integer; a float or a bool is refused
float	A JSON number, an integer included; a string or a bool is refused
datetime	ISO 8601 text
An enum	One of its values
tuple[T, ...], list[T]	A JSON array whose elements read as T
Mapping[str, T], dict[str, T]	A JSON object whose values read as T
A dataclass	A nested JSON object, read by the same rules
JsonValue	Any JSON value, unchecked below it

A missing field takes the dataclass default; a missing field without one is refused. A field the record does not declare is ignored, unless the record sets __pydantic_config__ = UNKNOWN_FIELDS_REFUSED, as settings and versioned files do, so a misspelled key is refused rather than dropped. Every refusal is collected into one pydantic.ValidationError, a ValueError whose errors() give each field's path as loc, such as ("metrics", "loss", "value").

The object is written back to JSON text and validated in pydantic's JSON mode, because strict Python-mode validation accepts a dataclass only as an instance. One validator per record type is built on first use and kept. Records are read at the edges of a program, never inside a function JAX traces.

substrax.records

Records read from JSON: frozen dataclasses validated by pydantic's strict JSON mode.

A record is a frozen dataclass; dump_record writes it as JSON values and read_record reads one back from the parsed JSON object. Each field is checked against its annotation as JSON defines the types: a number is never read from a string, a bool is not a number, an integer is accepted where a float is expected (JSON has one number type), nested dataclasses, datetime (ISO 8601 text), enums (by value) and tuples (from arrays) are read by their annotations, and a missing field takes its default. A refused record raises pydantic.ValidationError, a ValueError listing each field's path through the record.

The object is re-serialised to JSON text and validated in pydantic's JSON mode, because its strict Python mode reads a dataclass only from an instance of it; a record is small and read at the edge of a program, never inside a traced function.

UNKNOWN_FIELDS_REFUSED = ConfigDict(extra='forbid') module-attribute

Set as a record's __pydantic_config__ to refuse fields it does not declare, such as a misspelled setting; without it read_record ignores them.

read_record(record_type, data)

Read a record of record_type from a parsed JSON object.

Parameters:

Name	Type	Description	Default
record_type	type[T]	The record's dataclass.	required
data	Mapping[str, JsonValue]	The JSON object, as json.loads or Orbax's JsonRestore returns it.	required

Returns:

Type	Description
T	The record.

Raises:

Type	Description
ValidationError	If a field is missing or holds a value its annotation does not admit; errors() gives each field's path as loc.
TypeError	If the validator returns something other than the record type, which pydantic does not do for a dataclass.

dump_record(record)

Write a record as the JSON object read_record reads back into an equal record.

Nested records become objects, tuples arrays, enums their values, paths and datetime their text.

Parameters:

Name	Type	Description	Default
record	object	A dataclass instance.	required

Returns:

Type	Description
dict[str, JsonValue]	The JSON object, ready for json.dumps.

Raises:

Type	Description
TypeError	If record is not a dataclass instance.

substrax.typing

Type aliases shared across the avitai packages; importing this module imports no jax.

PyTree = Any module-attribute

Any JAX pytree: nested tuples, lists and dicts of arrays and leaves.

JsonValue = str | int | float | bool | list[JsonValue] | dict[str, JsonValue] | None

A value JSON can hold: a scalar, or a list or object of JSON values (as Orbax's tree types).

CheckpointState = dict[str, Any]

The state a :class:Checkpointable hands a checkpoint: arrays, plain-Python leaves and nested dictionaries, laid out by each implementer. Its values stay Any, the type Grain's DatasetIterator.get_state, PyTorch's Stateful.state_dict and flax's nnx.to_pure_dict give the same object: each implementer reads its own layout back.

Checkpointable

Bases: Protocol

An object a checkpoint restores: it hands over its state and takes it back.

The state is a dictionary a checkpoint store can write (arrays and plain-Python leaves), and set_state on an object built the same way continues exactly where the saved one stood. A training loop checkpoints its data iterator, callbacks and extensions through this protocol without importing their packages.

get_state()

The state to checkpoint.

set_state(state)

Take back a state get_state returned.

substrax.typing also holds CheckpointState, the dictionary a Checkpointable hands a checkpoint. substrax.nnx_typing holds the aliases that need jax or flax: NnxState, the state a module's nnx.split returns, and PathEntry, an entry of a jax key path; substrax.typing stays importable without jax.

substrax.nnx_typing

Type aliases over JAX and Flax NNX, shared across the avitai packages; this module imports both.

The aliases that import no jax are in :mod:substrax.typing.

PathEntry = jax.tree_util.DictKey | jax.tree_util.SequenceKey | jax.tree_util.GetAttrKey | jax.tree_util.FlattenedIndexKey module-attribute

An entry of a key path, as jax.tree.flatten_with_path and tree_map_with_path give it; a path is jax.tree_util.KeyPath[PathEntry]. A plain assignment, not a type statement: jax types the entry classes Any (jax/_src/tree_util.py), which a lazily evaluated alias turns into an unknown type under strict checking.

NnxState = nnx.State[Key, NnxState | nnx.Variable[jax.Array]]

A module's state, as nnx.split and nnx.state return it: keyed by attribute name or list index (flax's Key), each value a nested state or a Variable holding an array. A gradient taken with respect to a state has the same structure.