Overview

Diffract is a Python library for inspecting and tracking neural network weights. It is designed for workflows where you repeatedly compute metrics over model parameters and want the results to be reproducible, persisted, and easy to visualize.

When to use Diffract

Diffract shines when your analysis is iterative: multiple training runs, many metrics, or large models. If you just need one metric for one model once, a simple script may be faster.

What it solves:

  • No more one-off scripts — a single API (Session) for adding models, computing metrics, exporting, and plotting.

  • Fast iteration — parameters and computed fields persist (SQLite / HDF5 / Zarr) and are reused instead of recomputed after restarting. Reuse is keyed by field presence alone: changing a kernel’s configuration does not invalidate stored results — erase them explicitly to recompute (see Configuring kernels).

  • Consistent metrics — metrics are computed by kernels with declared dependencies; Diffract executes them in the correct order.

  • Scales to large models — chunked execution and configurable storage keep memory usage manageable.

Core concepts

Session

Session is the main entry point. It manages configuration, storage/cache resources, and provides the public API for all operations.

from diffract import Session

session = Session(profile="local")
with session:
    session.models.add(model, model_id="run-001")
    session.compute.apply("frob_norm", "stable_rank")
    results = session.results.export_metrics(
        "frob_norm", "stable_rank", export_format="dict"
    )

Parameters

When you call session.models.add(...), Diffract extracts model parameters and stores them. Each parameter is addressable by:

  • model_id — identifier you assign when adding a model

  • name — layer/parameter name from the model

  • uid — unique identifier auto-generated by Diffract

  • type — parameter type (dense, embedding, etc.)

Weight-slice schema boundary

Diffract analyzes weight slices — the stored weight matrices extracted from a model. Runtime activations and per-step gradients are outside this schema: they are not modeled, because identity carries a single categorical axis beyond the parameter name, the model_id.

To analyze one model across several checkpoints, add each checkpoint under its own model_id and compare them with session.filter(model_ids=[...]). Two limits follow from this single-axis identity:

  • model_ids sort lexicographically as strings, so a bare numeric suffix orders unexpectedly ("run-10" sorts before "run-2"). Zero-pad the suffix or keep an explicit ordering in the analysis layer; Diffract does not model a numeric step axis.

  • The characters @, [, ], ,, and / are reserved and may not appear in a model_id (see naming rules).

Fields and kernels

A field is a named value stored on each parameter (e.g., frob_norm, stable_rank).

A kernel is a function that produces one or more fields. Kernels declare their input fields (dependencies) and output fields. Diffract resolves the dependency graph and executes kernels in order.

Storage and cache

  • Storage persists parameters and computed fields across Python sessions.

  • Cache accelerates repeated reads of computed values.

Both are pluggable. Storage backends: RAM, SQLite, HDF5, Zarr, or hybrid. Cache backends: none, simple in-process LRU, or Redis.

What you can do

  • Add models from PyTorch, TensorFlow, Flax, ONNX, or plain NumPy array dicts (and custom extractors)

  • Compute 60+ built-in fields (norms, ranks, spectral metrics, heavy-tailed fits, etc.)

  • Configure kernel parameters at runtime

  • Filter computations by model, layer, or parameter type

  • Export results to dict, JSON, pandas, polars, or list

  • Render Plotly visualizations via the optional viz extra

Why Diffract vs ad-hoc scripts

Ad-hoc scripts

Diffract

Manage file paths and join keys manually

Structured data model with stable identifiers

Re-run everything after a change

Incremental: apply computes only fields not stored yet

Copy-paste metric code

Reusable kernels with automatic dependency resolution

Hardcoded storage

Configurable backends via profiles or config files

Next steps