Skip to content

Index

analysis

Unregistered analysis templates; register concrete jobs with @analysis(name).

AnalysisBase(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: BaseTrainer[C, M], Generic[C, M]

Trainer-backed analysis; subclasses implement selection and computation.

state: train_state.TrainState property writable

Observe replacement training states, including after donated JIT calls.

from_trainer(trainer: BaseTrainer[Any, M]) -> Self classmethod

Borrow the live trainer; never initialize state or open owned resources.

select(paths: list[str]) -> str | list[str] abstractmethod

Choose one path or a nonempty ordered list (the root path is '').

analyze(layer: nn.Module | list[nn.Module], inputs: DebugInputs | list[DebugInputs]) -> Any abstractmethod

Pure numerical computation; return a PyTree of arrays/scalars/None.

A string selection supplies a bound module and DebugInputs; a list supplies matching lists. Do not fetch batches, read self.state, render, convert tracers to NumPy, or mutate Python state here. Configuration may be read at trace time; parameters and inputs come from the supplied layer.

plot(results: Any) -> Figure | JSONValue abstractmethod

Render gathered CPU results on host zero; None suppresses the artifact.

AttentionHeatmapAnalysis(spec: ExecutionSpec, base: Optional[Node] = None)

Bases: AnalysisBase[C, M], Generic[C, M]

Unregistered analysis template; put it before your trainer in the bases.