xwm.tools

Checkpointing and model summaries.

Utilities: checkpointing and model introspection.

Modules:

Name	Description
checkpoint	Saving and loading models.
summary	Model introspection.

Functions:

Name	Description
load	Load parameters from path into a model with like's structure.
load_config	Read the config sidecar written by save.
load_state	Restore a TrainState.
save	Serialise model to path, optionally with a config sidecar.
save_state	Serialise a full TrainState (model, teacher, optimizer).
count_params	Total number of inexact-array scalars in tree.
param_bytes	Bytes occupied by the parameters, at their current dtypes.

load

load(path: str | Path, like: PyTree) -> PyTree

Load parameters from path into a model with like's structure.

Build like exactly as the saved model was built (same sizes, same keys are not required -- only the same shapes).

Source code in xwm/tools/checkpoint.py

def load(path: str | Path, like: PyTree) -> PyTree:
    """Load parameters from ``path`` into a model with ``like``'s structure.

    Build ``like`` exactly as the saved model was built (same sizes, same keys
    are not required -- only the same shapes).
    """
    return eqx.tree_deserialise_leaves(Path(path), like)

load_config

load_config(path: str | Path) -> dict[str, Any]

Read the config sidecar written by save.

Source code in xwm/tools/checkpoint.py

def load_config(path: str | Path) -> dict[str, Any]:
    """Read the config sidecar written by :func:`save`."""
    path = Path(path)
    sidecar = path.with_suffix(path.suffix + ".json")
    if not sidecar.exists():
        raise FileNotFoundError(f"no config sidecar at {sidecar}")
    return json.loads(sidecar.read_text())

load_state

load_state(path: str | Path, like: PyTree) -> PyTree

Restore a TrainState.

Build like with Trainer.init() on a freshly constructed model, which allocates the teacher and optimizer state with the right shapes.

Source code in xwm/tools/checkpoint.py

def load_state(path: str | Path, like: PyTree) -> PyTree:
    """Restore a :class:`~xwm.training.TrainState`.

    Build ``like`` with ``Trainer.init()`` on a freshly constructed model, which
    allocates the teacher and optimizer state with the right shapes.
    """
    return load(path, like)

save

save(path: str | Path, model: PyTree, *, config: dict[str, Any] | None = None) -> Path

Serialise model to path, optionally with a config sidecar.

Parameters:

Name	Type	Description	Default
path	str | Path	destination file; parent directories are created.	required
config	dict[str, Any] | None	JSON-serialisable constructor arguments, written to <path>.json so load has something to rebuild from.	None

Source code in xwm/tools/checkpoint.py

def save(path: str | Path, model: PyTree, *, config: dict[str, Any] | None = None) -> Path:
    """Serialise ``model`` to ``path``, optionally with a config sidecar.

    Args:
        path: destination file; parent directories are created.
        config: JSON-serialisable constructor arguments, written to
            ``<path>.json`` so :func:`load` has something to rebuild from.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    eqx.tree_serialise_leaves(path, model)
    if config is not None:
        path.with_suffix(path.suffix + ".json").write_text(json.dumps(config, indent=2))
    return path

save_state

save_state(path: str | Path, state: PyTree) -> Path

Serialise a full TrainState (model, teacher, optimizer).

Source code in xwm/tools/checkpoint.py

def save_state(path: str | Path, state: PyTree) -> Path:
    """Serialise a full :class:`~xwm.training.TrainState` (model, teacher, optimizer)."""
    return save(path, state)

count_params

count_params(tree: PyTree) -> int

Total number of inexact-array scalars in tree.

Source code in xwm/tools/summary.py

def count_params(tree: PyTree) -> int:
    """Total number of inexact-array scalars in ``tree``."""
    leaves = jax.tree_util.tree_leaves(eqx.filter(tree, eqx.is_inexact_array))
    return sum(int(x.size) for x in leaves)

param_bytes

param_bytes(tree: PyTree) -> int

Bytes occupied by the parameters, at their current dtypes.

Source code in xwm/tools/summary.py

def param_bytes(tree: PyTree) -> int:
    """Bytes occupied by the parameters, at their current dtypes."""
    leaves = jax.tree_util.tree_leaves(eqx.filter(tree, eqx.is_inexact_array))
    return sum(int(x.size) * x.dtype.itemsize for x in leaves)

summary

Note

xwm.tools.summary is both a submodule and the function it exports. The function is documented here under its canonical path; xwm.tools.summary(...) is the way to call it.

A parameter-count tree, in the spirit of torchinfo.

Parameters:

Name	Type	Description	Default
max_depth	int	how far to descend before summarising a subtree as a total.	2
collapse_lists	bool	print blocks[0] and note the repeat count instead of listing every identical transformer block.	True

Source code in xwm/tools/summary.py

def summary(model: PyTree, *, max_depth: int = 2, collapse_lists: bool = True) -> str:
    """A parameter-count tree, in the spirit of ``torchinfo``.

    Args:
        max_depth: how far to descend before summarising a subtree as a total.
        collapse_lists: print ``blocks[0]`` and note the repeat count instead of
            listing every identical transformer block.
    """
    total = count_params(model)
    lines = [
        f"{type(model).__name__}: {total:,} params ({param_bytes(model) / 1e6:.1f} MB)",
    ]

    def walk(module: Any, depth: int, prefix: str) -> None:
        if depth > max_depth:
            return
        children = _children(module)
        seen_prefixes: set[str] = set()
        for name, child in children:
            base = name.split("[")[0]
            if collapse_lists and "[" in name:
                if base in seen_prefixes:
                    continue
                seen_prefixes.add(base)
                repeats = sum(1 for n, _ in children if n.split("[")[0] == base)
                label = f"{base}[x{repeats}]"
                count = count_params(child) * repeats
            else:
                label, count = name, count_params(child)
            share = 100.0 * count / total if total else 0.0
            lines.append(
                f"{'  ' * depth}{prefix}{label}: {type(child).__name__} "
                f"-- {count:,} ({share:.1f}%)"
            )
            walk(child, depth + 1, prefix)

    walk(model, 1, "")
    return "\n".join(lines)