torchcell.viz.transformer_diagnostics.TransformerDiagnostics

class TransformerDiagnostics(base_dir: str, mplstyle_path: str | None = None)

Bases: object

Visualization class for transformer-specific diagnostic metrics.

save_and_log_figure(fig: Figure, name: str, timestamp_str: str | None = None) → None

Log figure to wandb only (no disk save).

Parameters:

• fig – Matplotlib figure

• name – Wandb panel key (e.g., “val_transformer_diagnostics/summary”)

• timestamp_str – Unused, kept for API compatibility

plot_attention_diagnostics(attention_stats: dict[int, dict[str, float]], residual_ratios: dict[int, float] | None = None, qk_logit_stats: dict[int, dict[str, float]] | None = None, gradient_norms: dict[int, float] | None = None, num_epochs: int = 0, stage: str = 'val') → None

Create aggregated diagnostic plot for transformer attention metrics.

Parameters:

• attention_stats – Dict mapping layer_idx -> { “entropy”: float, “effective_rank”: float, “top5”: float, “top10”: float, “top50”: float, “max_row_weight”: float, “col_entropy”: float, “max_col_sum”: float

• }

• residual_ratios – Optional dict mapping layer_idx -> residual_update_ratio

• qk_logit_stats – Optional dict mapping layer_idx -> { “logit_mean”: float, “logit_std”: float, “saturation_ratio”: float

• }

• gradient_norms – Optional dict mapping layer_idx -> gradient_norm (not currently plotted)

• num_epochs – Current epoch number

• stage – Stage name (e.g., “val”)