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), so you don’t recompute everything after restarting.

  • 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.)

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: only compute missing fields

Copy-paste metric code

Reusable kernels with automatic dependency resolution

Hardcoded storage

Configurable backends via profiles or config files

Next steps