base

Analysis jobs using the trainer's existing configuration and state lifecycle.

Use class MyAnalysis(AnalysisBase[C, M], MyTrainer). BaseTrainer owns config, construction, restoration/surgery, and cleanup; analysis replaces run(). Pass a checkpoint as base= to restore, or omit it to analyze initialized state. Execute through the normal job entry point in either case. Constructing an analysis does not call setup or initialize training state.

Implement select(paths), analyze(layer, inputs), and plot(results). Analysis returns a JAX-compatible PyTree; the framework gathers it before host zero calls plot(), which returns a Figure, JSON value, or None. Crop/reduce large arrays in analyze() before returning them. Analysis does not tick the trainer's node. select() may return a string or a nonempty list of strings. A list supplies parallel lists of bound modules and DebugInputs to analyze(), in selection order; even a one-item list stays a list. The first invocation of each path is captured inside one compiled trace. Notebook Debugger remains a separate eager interface. Each JSON/PDF is saved in results and as an artifact on the current node. Query with store.query().node(node).artifact().select(), or use get(node). Each analysis payload contains filename and content (the exact file bytes). Tandem analyses attach to the trainer's current node; NAME identifies the file and artifact, so distinct analyses on the same node must use distinct names.

For tandem use, declare ANALYSIS on the trainer. from_trainer() creates a borrowed view without construction or state initialization. Its node/state follow the trainer. Each run reads fresh state and prepares inputs outside the cached pure trace. The trainer calls run() at validation cadence; it alone owns setup, ticking, and cleanup. training/analyze disables this integration. Plotting helpers accept CPU data and never decide what to gather. CONFIG is the combined trainer and analysis dataclass schema; the trainer constructor hydrates it into self.args. Compose specialized schemas explicitly through dataclass inheritance.