Module neuralops.analysis

Analysis tools: plotting, metrics, layerwise analysis, model comparison, reporting.

Sub-modules

neuralops.analysis.comparison

Comparison utilities for analysis output.

neuralops.analysis.generate_all_modules_analysis

Comprehensive 217-Module Analysis Generator. Extracts and validates all 25 required columns for paper/_ALL_MODULES_ANALYSIS.csv and …

neuralops.analysis.generate_sweep_publication_artifacts

Publication-Grade Sweep Analysis, Reporting, and Visualization Engine …

neuralops.analysis.layerwise_analyzer

Streaming Layer-wise Analysis Engine for Foundation Models & Transformers …

neuralops.analysis.metrics

Metrics computation, aggregation, and tracking.

neuralops.analysis.model_comparison

Multi‑run model comparison and plotting.

neuralops.analysis.plotting

Comprehensive plotting functions for model analysis.

neuralops.analysis.report_generator

HTML, LaTeX, interactive report generators.

neuralops.analysis.utils

Utility functions for analysis.

Functions

def aggregate_results(results_list: list[dict[str, typing.Any]]) ‑> dict[str, typing.Any]

Expand source code

def aggregate_results(results_list: list[dict[str, Any]]) -> dict[str, Any]:
    """Aggregate results."""
    if not results_list:
        return {}
    aggregated = {}
    for key in results_list[0].keys():
        values = [r[key] for r in results_list if key in r and r[key] is not None]
        if values:
            aggregated[f"{key}_mean"] = np.mean(values)
            aggregated[f"{key}_std"] = np.std(values)
            aggregated[f"{key}_values"] = values
    return aggregated

Aggregate results.

def compute_metrics(predictions: numpy.ndarray, labels: numpy.ndarray, task: str = 'classification') ‑> dict[str, float]

Expand source code

def compute_metrics(
    predictions: np.ndarray, labels: np.ndarray, task: str = "classification"
) -> dict[str, float]:
    """Compute and return metrics."""
    from scipy.stats import pearsonr, spearmanr
    from sklearn.metrics import accuracy_score, f1_score, matthews_corrcoef

    metrics = {}
    if task in ("regression", "stsb"):
        pred_flat = np.squeeze(predictions)
        labels_flat = np.squeeze(labels)
        metrics["pearson"] = pearsonr(pred_flat, labels_flat)[0]
        metrics["spearman"] = spearmanr(pred_flat, labels_flat)[0]
        metrics["pearson_spearman"] = (metrics["pearson"] + metrics["spearman"]) / 2
    else:
        preds = np.argmax(predictions, axis=-1) if predictions.ndim > 1 else predictions
        metrics["accuracy"] = accuracy_score(labels, preds)
        metrics["f1_macro"] = f1_score(labels, preds, average="macro", zero_division=0)
        metrics["f1_weighted"] = f1_score(labels, preds, average="weighted", zero_division=0)
        if len(np.unique(labels)) == 2:
            metrics["matthews_corrcoef"] = matthews_corrcoef(labels, preds)
    return metrics

Compute and return metrics.

def correlation_matrix(df: pandas.DataFrame, columns: list[str]) ‑> pandas.DataFrame

Expand source code

def correlation_matrix(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    """Correlation matrix."""
    return df[columns].corr()

Correlation matrix.

def create_comprehensive_plots(df: pandas.DataFrame, output_dir: str, figsize: tuple[int, int] = (18, 16)) ‑> None

Expand source code

def create_comprehensive_plots(
    df: pd.DataFrame, output_dir: str, figsize: tuple[int, int] = (18, 16)
) -> None:
    """Build and return comprehensive plots."""
    os.makedirs(output_dir, exist_ok=True)
    set_style()
    fig, axes = plt.subplots(3, 3, figsize=figsize)
    fig.suptitle(
        "SVD Compression Analysis: Comprehensive Results",
        fontsize=16,
        fontweight="bold",
    )
    # 1. Accuracy vs Rank by model
    ax = axes[0, 0]
    for model in df["model"].unique():
        sub = df[df["model"] == model].groupby("rank")["accuracy"].mean().reset_index()
        ax.plot(sub["rank"], sub["accuracy"], marker="o", label=model, linewidth=2)
    ax.set_title("Accuracy vs Rank (by Model)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    # 2. Accuracy vs Rank by task
    ax = axes[0, 1]
    for task in df["task"].unique():
        sub = df[df["task"] == task].groupby("rank")["accuracy"].mean().reset_index()
        ax.plot(sub["rank"], sub["accuracy"], marker="s", label=task, linewidth=2)
    ax.set_title("Accuracy vs Rank (by Task)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    # 3. Speedup factor
    ax = axes[0, 2]
    if "speedup_factor" in df.columns:
        for model in df["model"].unique():
            sub = df[df["model"] == model].groupby("rank")["speedup_factor"].mean()
            ax.plot(sub.index, sub.values, marker="^", label=model, linewidth=2)
        ax.set_title("Speedup vs Rank")
        ax.legend()
        ax.grid(True, alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No speedup data",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
    # 4. Parameter reduction
    ax = axes[1, 0]
    if "param_reduction" in df.columns:
        pd_data = df.groupby(["model", "rank"])["param_reduction"].mean().reset_index()
        models = pd_data["model"].unique()
        width = 0.8 / len(models)
        for i, m in enumerate(models):
            sub = pd_data[pd_data["model"] == m]
            x = sub["rank"] + (i - len(models) / 2) * width
            ax.bar(x, sub["param_reduction"], width=width, label=m)
        ax.set_title("Parameter Reduction")
        ax.legend()
        ax.grid(True, axis="y", alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No param reduction",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
    # 5. Accuracy drop distribution
    ax = axes[1, 1]
    if "accuracy_drop" in df.columns:
        drops = df["accuracy_drop"].dropna()
        ax.hist(drops, bins=20, color="lightcoral", edgecolor="darkred", alpha=0.7)
        ax.axvline(drops.mean(), color="blue", linestyle="-", label=f"Mean: {drops.mean():.3f}")
        ax.set_title("Accuracy Drop Distribution")
        ax.legend()
        ax.grid(True, axis="y", alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No accuracy drop",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
    # 6. Memory reduction
    ax = axes[1, 2]
    if "memory_reduction" in df.columns:
        mem = df.groupby("rank")["memory_reduction"].mean().reset_index()
        ax.bar(
            mem["rank"],
            mem["memory_reduction"],
            color="lightgreen",
            edgecolor="darkgreen",
        )
        ax.set_title("Avg Memory Reduction")
        ax.grid(True, axis="y", alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No memory reduction",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
    # 7. Accuracy by task type
    ax = axes[2, 0]
    if "task_type" in df.columns:
        for tt in df["task_type"].unique():
            sub = df[df["task_type"] == tt].groupby("rank")["accuracy"].mean().reset_index()
            ax.plot(sub["rank"], sub["accuracy"], marker="D", label=tt, linewidth=2)
        ax.set_title("Accuracy by Task Type")
        ax.legend()
        ax.grid(True, alpha=0.3)
    else:
        ax.text(0.5, 0.5, "No task type", ha="center", va="center", transform=ax.transAxes)
    # 8. Compression ratio heatmap
    ax = axes[2, 1]
    if "compression_ratio" in df.columns:
        pivot = df.pivot_table(
            index="model", columns="rank", values="compression_ratio", aggfunc="mean"
        )
        im = ax.imshow(pivot.values, cmap="YlOrRd", aspect="auto")
        ax.set_xticks(range(len(pivot.columns)))
        ax.set_xticklabels(pivot.columns)
        ax.set_yticks(range(len(pivot.index)))
        ax.set_yticklabels(pivot.index)
        plt.colorbar(im, ax=ax, label="Compression Ratio")
        ax.set_title("Compression Ratio Heatmap")
    else:
        ax.text(0.5, 0.5, "No comp ratio", ha="center", va="center", transform=ax.transAxes)
    # 9. Best config per dataset
    ax = axes[2, 2]
    best = []
    for task in df["task"].unique():
        sub = df[df["task"] == task]
        if "accuracy" in sub.columns:
            idx = sub["accuracy"].idxmax()
            best.append(sub.loc[idx])
    if best:
        best_df = pd.DataFrame(best)
        ax.barh(best_df["task"], best_df["accuracy"], color="skyblue", edgecolor="navy")
        ax.set_title("Best Configuration per Dataset")
        for i, (acc, model, rank) in enumerate(
            zip(best_df["accuracy"], best_df["model"], best_df["rank"])
        ):
            ax.text(acc + 0.005, i, f"{model} r={int(rank)}", va="center", fontsize=8)
    else:
        ax.text(0.5, 0.5, "No best config", ha="center", va="center", transform=ax.transAxes)
    plt.tight_layout()
    plt.savefig(Path(output_dir) / "comprehensive_analysis.png", dpi=300, bbox_inches="tight")
    plt.close()

Build and return comprehensive plots.

def detect_least_compressible_layers(model, tokenizer, calibration_texts: list[str], k: int = 5) ‑> list[int]

Expand source code

def detect_least_compressible_layers(
    model, tokenizer, calibration_texts: list[str], k: int = 5
) -> list[int]:
    """Detect least compressible layers."""
    from .layerwise_analyzer import LayerwiseAnalyzer

    analyzer = LayerwiseAnalyzer(model, tokenizer)
    result = analyzer.analyze(calibration_texts)
    changes = result["cosine_similarities"]
    most_change_scores = [1 - sim for sim in changes]
    top_indices = np.argsort(most_change_scores)[-k:]
    return sorted(top_indices)

Detect least compressible layers.

def estimate_intrinsic_dimension(hidden_states: np.ndarray, k: int = 5) ‑> float

Expand source code

def estimate_intrinsic_dimension(hidden_states: np.ndarray, k: int = 5) -> float:
    """
    Estimates the intrinsic dimension of a representation manifold using k-NN distances.

    Args:
        hidden_states: Sample representation array of shape [N, D].
        k: Neighbor rank for distance estimation.

    Returns:
        float: Estimated intrinsic dimension between 1.0 and D.
    """
    hidden_states = np.asarray(hidden_states, dtype=np.float64)
    if hidden_states.ndim == 1:
        return 1.0
    N, D = hidden_states.shape
    if k + 1 > N:
        return float(min(N, D))

    from scipy.spatial.distance import pdist, squareform

    try:
        dists = squareform(pdist(hidden_states))
        np.fill_diagonal(dists, np.inf)
        kth_dists = np.sort(dists, axis=1)[:, k - 1]
        valid = kth_dists > 1e-10
        if not valid.any():
            return float(D)

        # Maximum likelihood estimation of intrinsic dimension
        log_ratios = np.log(kth_dists[valid] / (kth_dists[valid][:, None] + 1e-12))
        id_estimate = -1.0 / (np.mean(log_ratios) + 1e-12)
        return float(max(1.0, min(float(D), abs(id_estimate))))
    except Exception:
        return float(min(N, D))

Estimates the intrinsic dimension of a representation manifold using k-NN distances.

Args

-----=

hidden_states

Sample representation array of shape [N, D].

k

Neighbor rank for distance estimation.

Returns

-----=

float

Estimated intrinsic dimension between 1.0 and D.

def filter_results(df: pandas.DataFrame,
models: list[str] | None = None,
tasks: list[str] | None = None,
ranks: list[int] | None = None,
min_accuracy: float | None = None) ‑> pandas.DataFrame

Expand source code

def filter_results(
    df: pd.DataFrame,
    models: list[str] | None = None,
    tasks: list[str] | None = None,
    ranks: list[int] | None = None,
    min_accuracy: float | None = None,
) -> pd.DataFrame:
    """Filter results."""
    result = df.copy()
    if models:
        result = result[result["model"].isin(models)]
    if tasks:
        result = result[result["task"].isin(tasks)]
    if ranks:
        result = result[result["rank"].isin(ranks)]
    if min_accuracy is not None and "accuracy" in result.columns:
        result = result[result["accuracy"] >= min_accuracy]
    return result

Filter results.

def generate_full_report(df: pandas.DataFrame, output_dir: str, experiment_name: str = 'experiment') ‑> str

Expand source code

def generate_full_report(
    df: pd.DataFrame, output_dir: str, experiment_name: str = "experiment"
) -> str:
    """Generate full report."""
    os.makedirs(output_dir, exist_ok=True)
    plots_dir = os.path.join(output_dir, "plots")
    from .plotting import create_comprehensive_plots, plot_tradeoff

    create_comprehensive_plots(df, plots_dir)
    plot_tradeoff(df, plots_dir)
    generate_html_report(df, output_dir, title=f"{experiment_name} Full Report")
    csv_path = os.path.join(output_dir, "summary.csv")
    df.to_csv(csv_path, index=False)
    generate_latex_table(df, os.path.join(output_dir, "results.tex"))
    return os.path.join(output_dir, "report.html")

Generate full report.

def generate_html_report(df: pandas.DataFrame, output_dir: str, title: str = 'Analysis Report') ‑> str

Expand source code

def generate_html_report(df: pd.DataFrame, output_dir: str, title: str = "Analysis Report") -> str:
    """Generate html report."""
    os.makedirs(output_dir, exist_ok=True)
    html = f"""<!DOCTYPE html>
<html><head><meta charset='utf-8'><title>{title}</title>
<style>body{{font-family:Arial;margin:20px}} table{{border-collapse:collapse;width:100%}} th,td{{border:1px solid #ddd;padding:8px}} th{{background:#f2f2f2}}</style>
</head><body>
<h1>{title}</h1><p>Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</p>
<h2>Summary Table</h2>{df.to_html(index=False)}
</body></html>"""
    out = os.path.join(output_dir, "report.html")
    with open(out, "w") as f:
        f.write(html)
    return out

Generate html report.

def generate_interactive_report(df: pandas.DataFrame, output_dir: str) ‑> str

Expand source code

def generate_interactive_report(df: pd.DataFrame, output_dir: str) -> str:
    """Generate interactive report."""
    os.makedirs(output_dir, exist_ok=True)
    import plotly.graph_objects as go

    fig1 = go.Figure()
    for task in df["task"].unique():
        sub = df[df["task"] == task]
        fig1.add_trace(
            go.Scatter(x=sub["rank"], y=sub["accuracy"], mode="lines+markers", name=task)
        )
    fig1.update_layout(title="Accuracy vs Rank", xaxis_title="Rank", yaxis_title="Accuracy")
    fig2 = go.Figure()
    for model in df["model"].unique():
        sub = df[df["model"] == model]
        fig2.add_trace(
            go.Scatter(x=sub["rank"], y=sub["accuracy"], mode="lines+markers", name=model)
        )
    fig2.update_layout(title="Model Comparison", xaxis_title="Rank", yaxis_title="Accuracy")
    html = f"""<!DOCTYPE html>
<html><head><script src="https://cdn.plot.ly/plotly-latest.min.js"></script></head>
<body><h1>Interactive Report</h1>
<div id="rankPlot"></div><div id="modelPlot"></div>
<script>
Plotly.newPlot('rankPlot', {json.dumps(fig1.to_dict())});
Plotly.newPlot('modelPlot', {json.dumps(fig2.to_dict())});
</script></body></html>"""
    path = os.path.join(output_dir, "interactive_report.html")
    with open(path, "w") as f:
        f.write(html)
    return path

Generate interactive report.

def generate_latex_table(df: pandas.DataFrame, output_path: str, caption: str = 'Results') ‑> str

Expand source code

def generate_latex_table(df: pd.DataFrame, output_path: str, caption: str = "Results") -> str:
    """Generate latex table."""
    lines = [
        r"\begin{table}[htbp]\centering",
        r"\caption{" + caption + "}",
        r"\begin{tabular}{lcccr}",
        r"\toprule",
    ]
    lines.append("Model & Task & Rank & Accuracy & Drop \\\\")
    lines.append(r"\midrule")
    for _, row in df.iterrows():
        acc = row.get("accuracy", 0)
        drop = row.get("accuracy_drop", "-")
        drop_str = f"{drop:.4f}" if isinstance(drop, (int, float)) else "-"
        lines.append(
            f"{row['model']} & {row['task']} & {row['rank']} & {acc:.4f} & {drop_str} \\\\"
        )
    lines.extend([r"\bottomrule", r"\end{tabular}", r"\end{table}"])
    with open(output_path, "w") as f:
        f.write("\n".join(lines))
    return output_path

Generate latex table.

def get_best_configurations(df: pandas.DataFrame, metric: str = 'accuracy') ‑> pandas.DataFrame

Expand source code

def get_best_configurations(df: pd.DataFrame, metric: str = "accuracy") -> pd.DataFrame:
    """Return the best configurations."""
    best_rows = []
    for (task, model), group in df.groupby(["task", "model"]):
        if metric in group.columns:
            idx = group[metric].idxmax()
            best_rows.append(group.loc[idx])
    return pd.DataFrame(best_rows)

Return the best configurations.

def plot_accuracy_curve(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_accuracy_curve(history, metrics, config, output_dir, opts) -> Any:
    """Plot accuracy curve."""
    steps = [h.get("step", i) for i, h in enumerate(history) if "eval_accuracy" in h]
    acc = [h["eval_accuracy"] for h in history if "eval_accuracy" in h]
    if not steps:
        return []
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 6]))
    ax.plot(steps, acc, "o-", color="#5DC963", markersize=4)
    ax.set_xlabel("Step")
    ax.set_ylabel("Accuracy")
    ax.set_title("Validation Accuracy")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "accuracy_curve.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot accuracy curve.

def plot_attention_heatmap(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_attention_heatmap(history, metrics, config, output_dir, opts) -> Any:
    """Plot attention heatmap."""
    attn_path = Path(output_dir).parent / "artifacts" / "attention_sample.json"
    if not attn_path.exists():
        return []
    with open(attn_path) as f:
        attn_data = json.load(f)
    matrix = np.array(attn_data.get("matrix", []))
    tokens = attn_data.get("tokens", None)
    layer = attn_data.get("layer", 0)
    head = attn_data.get("head", 0)
    # Simple heatmap
    fig, ax = plt.subplots(figsize=(10, 8))
    sns.heatmap(matrix, cmap="viridis", xticklabels=tokens, yticklabels=tokens, ax=ax)
    ax.set_title(f"Attention Heatmap – Layer {layer} Head {head}")
    path = Path(output_dir) / "attention_heatmap.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot attention heatmap.

def plot_confusion_matrix(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_confusion_matrix(history, metrics, config, output_dir, opts) -> Any:
    """Plot confusion matrix."""
    preds_path = Path(output_dir).parent / "artifacts" / "eval_predictions.json"
    if not preds_path.exists():
        return []
    with open(preds_path) as f:
        data = json.load(f)
    labels = data.get("labels", [])
    preds = data.get("predictions", [])
    if not labels:
        return []
    from sklearn.metrics import confusion_matrix

    cm = confusion_matrix(labels, preds)
    class_names = opts.get("class_names", None)
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 8]))
    sns.heatmap(
        cm,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=class_names,
        yticklabels=class_names,
        ax=ax,
    )
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title("Confusion Matrix")
    path = Path(output_dir) / "confusion_matrix.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot confusion matrix.

def plot_gpu_memory(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_gpu_memory(history, metrics, config, output_dir, opts) -> Any:
    """Plot gpu memory."""
    gpu_steps = [h.get("step", i) for i, h in enumerate(history) if "gpu_mem" in h]
    gpu = [h["gpu_mem"] for h in history if "gpu_mem" in h]
    cpu_steps = [h.get("step", i) for i, h in enumerate(history) if "cpu_mem" in h]
    cpu = [h["cpu_mem"] for h in history if "cpu_mem" in h]
    if not gpu_steps and not cpu_steps:
        return []
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 4]))
    if gpu_steps:
        ax.plot(gpu_steps, gpu, label="GPU Mem (MB)", color="#3B528B")
    if cpu_steps:
        ax.plot(cpu_steps, cpu, label="CPU Mem (MB)", color="#5DC963")
    ax.set_xlabel("Step")
    ax.set_ylabel("Memory (MB)")
    ax.set_title("Memory Usage")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "gpu_memory.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot gpu memory.

def plot_gradient_flow(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_training_time(history, metrics, config, output_dir, opts) -> Any:
    """Plot training time."""
    times = [
        h.get("elapsed", i * 0.01) for i, h in enumerate(history) if "elapsed" in h or "step" in h
    ]
    if not times:
        return []
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 4]))
    ax.plot(range(len(times)), times, color="#440154")
    ax.set_xlabel("Step")
    ax.set_ylabel("Time (s)")
    ax.set_title("Training Time vs Steps")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "training_time.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot training time.

def plot_loss_curve(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_loss_curve(history, metrics, config, output_dir, opts) -> Any:
    """Plot loss curve."""
    steps = [h.get("step", i) for i, h in enumerate(history) if "loss" in h]
    train = [h["loss"] for h in history if "loss" in h]
    eval_steps = [h.get("step", i) for i, h in enumerate(history) if "eval_loss" in h]
    eval_l = [h["eval_loss"] for h in history if "eval_loss" in h]

    smooth = opts.get("loss_smooth", 0.6)
    if smooth > 0 and len(train) > 5:
        window = max(1, int(len(train) * smooth / 10))
        train = uniform_filter1d(np.array(train), size=window)

    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 6]))
    ax.plot(steps, train, label="Train Loss", color="#3B528B")
    if eval_steps:
        ax.plot(eval_steps, eval_l, "o-", label="Eval Loss", color="#21908C", markersize=4)
    ax.set_xlabel("Step")
    ax.set_ylabel("Loss")
    ax.set_title("Training & Validation Loss")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "loss_curve.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot loss curve.

def plot_lr_schedule(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_lr_schedule(history, metrics, config, output_dir, opts) -> Any:
    """Plot lr schedule."""
    steps = [h.get("step", i) for i, h in enumerate(history) if "learning_rate" in h]
    lr = [h["learning_rate"] for h in history if "learning_rate" in h]
    if not steps:
        return []
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 4]))
    ax.plot(steps, lr, color="#FDE725")
    ax.set_xlabel("Step")
    ax.set_ylabel("LR")
    ax.set_title("Learning Rate Schedule")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "lr_schedule.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot lr schedule.

def plot_parameter_distribution(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_parameter_distribution(history, metrics, config, output_dir, opts) -> Any:
    """Plot parameter distribution."""
    param_path = Path(output_dir).parent / "artifacts" / "param_counts.json"
    if not param_path.exists():
        return []
    with open(param_path) as f:
        param_data = json.load(f)
    layers = list(param_data.keys())
    counts = list(param_data.values())
    fig, ax = plt.subplots(figsize=opts.get("figsize", [8, 6]))
    ax.barh(layers, counts, color="#21908C")
    ax.set_xlabel("Parameter Count")
    ax.set_title("Parameter Distribution per Layer")
    plt.tight_layout()
    path = Path(output_dir) / "parameter_count.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot parameter distribution.

def plot_pareto_front(df, output_dir, x_metric='speedup_factor', y_metric='accuracy') ‑> None

Expand source code

def plot_pareto_front(df, output_dir, x_metric="speedup_factor", y_metric="accuracy") -> None:
    """Plot pareto front."""
    set_style()
    points = df[[x_metric, y_metric, "model", "rank"]].dropna().values
    pareto = []
    for i, p in enumerate(points):
        dominated = False
        for q in points:
            if q[0] >= p[0] and q[1] >= p[1] and (q[0] > p[0] or q[1] > p[1]):
                dominated = True
                break
        if not dominated:
            pareto.append(p)
    pareto_df = pd.DataFrame(pareto, columns=[x_metric, y_metric, "model", "rank"])
    fig, ax = plt.subplots(figsize=(10, 8))
    for model in df["model"].unique():
        sub = df[df["model"] == model]
        ax.scatter(sub[x_metric], sub[y_metric], label=model, alpha=0.5)
    ax.scatter(
        pareto_df[x_metric],
        pareto_df[y_metric],
        color="red",
        s=150,
        marker="*",
        label="Pareto front",
    )
    ax.set_xlabel(x_metric.replace("_", " ").capitalize())
    ax.set_ylabel(y_metric.capitalize())
    ax.set_title("Pareto Frontier of Configurations")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(Path(output_dir) / "pareto_front.png", dpi=150)
    plt.close()

Plot pareto front.

def plot_radar_chart(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_radar_chart(history, metrics, config, output_dir, opts) -> Any:
    """Plot radar chart."""
    if not metrics:
        return []
    numeric = {k: v for k, v in metrics.items() if isinstance(v, (int, float))}
    if len(numeric) < 3:
        return []
    categories = list(numeric.keys())
    values = list(numeric.values())
    angles = np.linspace(0, 2 * np.pi, len(categories), endpoint=False).tolist()
    values = values + values[:1]
    angles += angles[:1]
    fig, ax = plt.subplots(figsize=opts.get("figsize", [8, 8]), subplot_kw=dict(polar=True))
    ax.fill(angles, values, alpha=0.25, color="#3B528B")
    ax.plot(angles, values, color="#3B528B", linewidth=2)
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(categories)
    ax.set_title("Performance Radar Chart")
    plt.tight_layout()
    path = Path(output_dir) / "radar_chart.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot radar chart.

def plot_rank_allocation(recommended_ranks, output_dir) ‑> None

Expand source code

def plot_rank_allocation(recommended_ranks, output_dir) -> None:
    """Plot rank allocation."""
    set_style()
    fig, ax = plt.subplots(figsize=(12, 6))
    layers = list(range(len(recommended_ranks)))
    ax.bar(layers, recommended_ranks, color="steelblue", edgecolor="navy")
    ax.set_xlabel("Layer Index")
    ax.set_ylabel("Recommended Rank")
    ax.set_title("Layer‑wise Rank Allocation Recommendation")
    ax.grid(True, axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(Path(output_dir) / "rank_allocation.png", dpi=150)
    plt.close()

Plot rank allocation.

def plot_tradeoff(df, output_dir, metric='accuracy') ‑> None

Expand source code

def plot_tradeoff(df, output_dir, metric="accuracy") -> None:
    """Plot tradeoff."""
    set_style()
    fig, ax = plt.subplots(figsize=(10, 8))
    for model in df["model"].unique():
        sub = df[df["model"] == model].dropna(subset=["speedup_factor", metric])
        ax.scatter(sub["speedup_factor"], sub[metric], s=100, label=model, alpha=0.7)
    ax.set_xlabel("Speedup Factor")
    ax.set_ylabel(metric.capitalize())
    ax.set_title("Accuracy vs Speedup Tradeoff")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(Path(output_dir) / "tradeoff.png", dpi=150, bbox_inches="tight")
    plt.close()

Plot tradeoff.

def plot_training_history(history: dict, save_path: str) ‑> None

Expand source code

def plot_training_history(history: dict, save_path: str) -> None:
    """Plot training history."""
    fig, ax = plt.subplots(figsize=(8, 5))
    if history.get("train_steps") and history.get("train_loss"):
        ax.plot(
            history["train_steps"],
            history["train_loss"],
            label="Train Loss",
            color="blue",
        )
    if history.get("eval_steps") and history.get("eval_loss"):
        ax.plot(
            history["eval_steps"],
            history["eval_loss"],
            label="Eval Loss",
            color="red",
            marker="o",
        )
    ax.set_xlabel("Steps")
    ax.set_ylabel("Loss")
    ax.set_title("Training History")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()

Plot training history.

def plot_training_time(history, metrics, config, output_dir, opts) ‑> Any

Expand source code

def plot_training_time(history, metrics, config, output_dir, opts) -> Any:
    """Plot training time."""
    times = [
        h.get("elapsed", i * 0.01) for i, h in enumerate(history) if "elapsed" in h or "step" in h
    ]
    if not times:
        return []
    fig, ax = plt.subplots(figsize=opts.get("figsize", [10, 4]))
    ax.plot(range(len(times)), times, color="#440154")
    ax.set_xlabel("Step")
    ax.set_ylabel("Time (s)")
    ax.set_title("Training Time vs Steps")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    path = Path(output_dir) / "training_time.png"
    fig.savefig(path, dpi=opts.get("dpi", 150), bbox_inches="tight")
    plt.close()
    return [str(path)]

Plot training time.

def quick_compare(results_dirs: list[str], output_dir: str = './results/comparison') ‑> None

Expand source code

def quick_compare(results_dirs: list[str], output_dir: str = "./results/comparison") -> None:
    """Quick compare."""
    comparator = ModelComparison(results_dirs, output_dir)
    comparator.load_all_results()
    comparator.generate_all_plots()

Quick compare.

def set_style() ‑> None

Expand source code

def set_style() -> None:
    """Set the style."""
    plt.style.use("seaborn-v0_8-darkgrid")
    sns.set_palette("husl")
    sns.set_context("paper", font_scale=1.2)

Set the style.

def statistical_tests(df: pandas.DataFrame, group_col: str, value_col: str) ‑> dict[str, typing.Any]

Expand source code

def statistical_tests(df: pd.DataFrame, group_col: str, value_col: str) -> dict[str, Any]:
    """Statistical tests."""
    from scipy import stats

    groups = [g[1][value_col].dropna() for g in df.groupby(group_col)]
    if len(groups) < 2:
        return {"error": "Need at least two groups"}
    f_stat, p_anova = stats.f_oneway(*groups)
    pairwise = {}
    for i in range(len(groups) - 1):
        t, p = stats.ttest_ind(groups[i], groups[i + 1])
        pairwise[f"{i}_vs_{i + 1}"] = {"t": t, "p": p}
    return {"anova_p": p_anova, "pairwise": pairwise}

Statistical tests.

Classes

class LayerwiseAnalyzer (model: nn.Module, tokenizer: Any, device: str | torch.device | None = None)

Expand source code

class LayerwiseAnalyzer:
    """
    Streaming Layerwise Analyzer that captures inter-layer representations,
    outlier distribution spikes, and effective rank without accumulating VRAM memory.
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any,
        device: str | torch.device | None = None,
    ) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.device = str(device) if device is not None else self._get_device()
        self.model.to(self.device)
        self.model.eval()
        self.layer_accumulators: list[StreamingLayerAccumulator] = []
        self.hooks: list[torch.utils.hooks.RemovableHandle] = []

    def _get_device(self) -> str:
        """Detect best available compute device."""
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        return "cpu"

    def _find_transformer_layers(self) -> list[tuple[str, nn.Module]]:
        """Identify candidate Transformer / Encoder / Decoder block modules."""
        target_layers = []
        for name, module in self.model.named_modules():
            name_lower = name.lower()
            # Match standard transformer layer patterns
            if any(
                pattern in name_lower
                for pattern in [
                    "encoder.layer.",
                    "transformer.layer.",
                    "transformer.h.",
                    "model.layers.",
                    "layers.",
                    "h.",
                    "decoder.layer.",
                    "transformer.blocks.",
                ]
            ):
                # Filter to only the top-level block (e.g. layers.0, not layers.0.attention)
                parts = name.split(".")
                if parts[-1].isdigit():
                    target_layers.append((name, module))

        # Fallback: if no digit-indexed layers found, look for any layer module
        if not target_layers:
            for name, module in self.model.named_modules():
                if "layer" in name.lower() and isinstance(module, (nn.ModuleList, nn.Sequential)):
                    for sub_idx, sub_mod in enumerate(module):
                        target_layers.append((f"{name}.{sub_idx}", sub_mod))
                    break
        return target_layers

    def register_streaming_hooks(self) -> None:
        """Registers streaming forward hooks on all model transformer layers."""
        self.remove_hooks()
        layers = self._find_transformer_layers()

        num_layers = len(layers)
        self.layer_accumulators = [StreamingLayerAccumulator() for _ in range(num_layers)]

        for layer_idx, (layer_name, module) in enumerate(layers):

            def make_hook(idx: int):
                def hook_fn(mod: nn.Module, inputs: Any, output: Any) -> None:
                    # Extract tensor
                    tensor = output[0] if isinstance(output, tuple) else output
                    if isinstance(tensor, torch.Tensor):
                        # FIX: Update streaming accumulator dynamically and release activation immediately
                        self.layer_accumulators[idx].update(tensor)
                        # Explicitly release references
                        del tensor

                return hook_fn

            hook_handle = module.register_forward_hook(make_hook(layer_idx))
            self.hooks.append(hook_handle)

    def remove_hooks(self) -> None:
        """Remove all active hooks and free memory."""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

    def analyze(self, texts: list[str], batch_size: int = 8) -> dict[str, Any]:
        """
        Executes streaming layerwise analysis on input text corpus.

        Args:
            texts: List of input string prompts/sentences.
            batch_size: Batch size for tokenization and model inference.

        Returns:
            Dict containing cosine similarities, effective ranks, outlier statistics,
            and least-compressible layer rankings.
        """
        self.register_streaming_hooks()

        try:
            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i : i + batch_size]
                encoded = self.tokenizer(
                    batch_texts,
                    return_tensors="pt",
                    truncation=True,
                    max_length=128,
                    padding=True,
                )
                inputs = {k: v.to(self.device) for k, v in encoded.items()}
                with torch.no_grad():
                    self.model(**inputs)

                # Cleanup batch inputs
                del inputs, encoded
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        finally:
            self.remove_hooks()

        num_layers = len(self.layer_accumulators)
        if num_layers == 0:
            return {
                "num_layers": 0,
                "cosine_similarities": [],
                "effective_ranks": [],
                "max_outlier_magnitudes": [],
                "outlier_ratios": [],
                "kurtosis_scores": [],
                "intrinsic_dimensions": [],
            }

        # Extract per-layer mean vectors
        mean_vectors = [
            (
                acc.running_mean_vector
                if acc.running_mean_vector is not None
                else np.zeros(128, dtype=np.float32)
            )
            for acc in self.layer_accumulators
        ]

        # Compute cosine similarity between adjacent layers
        similarities = []
        for i in range(num_layers - 1):
            v1, v2 = mean_vectors[i], mean_vectors[i + 1]
            norm1, norm2 = np.linalg.norm(v1), np.linalg.norm(v2)
            if norm1 > 1e-8 and norm2 > 1e-8:
                sim = 1.0 - float(cosine(v1, v2))
                sim_clamped = max(0.0, min(1.0, sim))
            else:
                sim_clamped = 1.0
            similarities.append(sim_clamped)

        effective_ranks = [acc.compute_effective_rank() for acc in self.layer_accumulators]
        max_outliers = [acc.max_activation for acc in self.layer_accumulators]
        outlier_ratios = [acc.compute_outlier_ratio() for acc in self.layer_accumulators]
        kurtosis_scores = [acc.compute_kurtosis() for acc in self.layer_accumulators]

        # Estimate intrinsic dimension using mean vectors
        np.stack(mean_vectors) if num_layers > 1 else np.zeros((1, 128))
        intrinsic_dims = [
            estimate_intrinsic_dimension(
                acc.cov_matrix if acc.cov_matrix is not None else np.eye(10), k=3
            )
            for acc in self.layer_accumulators
        ]

        return {
            "num_layers": num_layers,
            "cosine_similarities": similarities,
            "effective_ranks": effective_ranks,
            "max_outlier_magnitudes": max_outliers,
            "outlier_ratios": outlier_ratios,
            "kurtosis_scores": kurtosis_scores,
            "intrinsic_dimensions": intrinsic_dims,
        }

Streaming Layerwise Analyzer that captures inter-layer representations, outlier distribution spikes, and effective rank without accumulating VRAM memory.

Methods

def analyze(self, texts: list[str], batch_size: int = 8) ‑> dict[str, typing.Any]

Expand source code

def analyze(self, texts: list[str], batch_size: int = 8) -> dict[str, Any]:
    """
    Executes streaming layerwise analysis on input text corpus.

    Args:
        texts: List of input string prompts/sentences.
        batch_size: Batch size for tokenization and model inference.

    Returns:
        Dict containing cosine similarities, effective ranks, outlier statistics,
        and least-compressible layer rankings.
    """
    self.register_streaming_hooks()

    try:
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i : i + batch_size]
            encoded = self.tokenizer(
                batch_texts,
                return_tensors="pt",
                truncation=True,
                max_length=128,
                padding=True,
            )
            inputs = {k: v.to(self.device) for k, v in encoded.items()}
            with torch.no_grad():
                self.model(**inputs)

            # Cleanup batch inputs
            del inputs, encoded
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    finally:
        self.remove_hooks()

    num_layers = len(self.layer_accumulators)
    if num_layers == 0:
        return {
            "num_layers": 0,
            "cosine_similarities": [],
            "effective_ranks": [],
            "max_outlier_magnitudes": [],
            "outlier_ratios": [],
            "kurtosis_scores": [],
            "intrinsic_dimensions": [],
        }

    # Extract per-layer mean vectors
    mean_vectors = [
        (
            acc.running_mean_vector
            if acc.running_mean_vector is not None
            else np.zeros(128, dtype=np.float32)
        )
        for acc in self.layer_accumulators
    ]

    # Compute cosine similarity between adjacent layers
    similarities = []
    for i in range(num_layers - 1):
        v1, v2 = mean_vectors[i], mean_vectors[i + 1]
        norm1, norm2 = np.linalg.norm(v1), np.linalg.norm(v2)
        if norm1 > 1e-8 and norm2 > 1e-8:
            sim = 1.0 - float(cosine(v1, v2))
            sim_clamped = max(0.0, min(1.0, sim))
        else:
            sim_clamped = 1.0
        similarities.append(sim_clamped)

    effective_ranks = [acc.compute_effective_rank() for acc in self.layer_accumulators]
    max_outliers = [acc.max_activation for acc in self.layer_accumulators]
    outlier_ratios = [acc.compute_outlier_ratio() for acc in self.layer_accumulators]
    kurtosis_scores = [acc.compute_kurtosis() for acc in self.layer_accumulators]

    # Estimate intrinsic dimension using mean vectors
    np.stack(mean_vectors) if num_layers > 1 else np.zeros((1, 128))
    intrinsic_dims = [
        estimate_intrinsic_dimension(
            acc.cov_matrix if acc.cov_matrix is not None else np.eye(10), k=3
        )
        for acc in self.layer_accumulators
    ]

    return {
        "num_layers": num_layers,
        "cosine_similarities": similarities,
        "effective_ranks": effective_ranks,
        "max_outlier_magnitudes": max_outliers,
        "outlier_ratios": outlier_ratios,
        "kurtosis_scores": kurtosis_scores,
        "intrinsic_dimensions": intrinsic_dims,
    }

Executes streaming layerwise analysis on input text corpus.

Args

-----=

texts

List of input string prompts/sentences.

batch_size

Batch size for tokenization and model inference.

Returns -----= Dict containing cosine similarities, effective ranks, outlier statistics, and least-compressible layer rankings.

def register_streaming_hooks(self) ‑> None

Expand source code

def register_streaming_hooks(self) -> None:
    """Registers streaming forward hooks on all model transformer layers."""
    self.remove_hooks()
    layers = self._find_transformer_layers()

    num_layers = len(layers)
    self.layer_accumulators = [StreamingLayerAccumulator() for _ in range(num_layers)]

    for layer_idx, (layer_name, module) in enumerate(layers):

        def make_hook(idx: int):
            def hook_fn(mod: nn.Module, inputs: Any, output: Any) -> None:
                # Extract tensor
                tensor = output[0] if isinstance(output, tuple) else output
                if isinstance(tensor, torch.Tensor):
                    # FIX: Update streaming accumulator dynamically and release activation immediately
                    self.layer_accumulators[idx].update(tensor)
                    # Explicitly release references
                    del tensor

            return hook_fn

        hook_handle = module.register_forward_hook(make_hook(layer_idx))
        self.hooks.append(hook_handle)

Registers streaming forward hooks on all model transformer layers.

def remove_hooks(self) ‑> None

Expand source code

def remove_hooks(self) -> None:
    """Remove all active hooks and free memory."""
    for hook in self.hooks:
        hook.remove()
    self.hooks.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

Remove all active hooks and free memory.

class MetricsTracker (filepath: str)

Expand source code

class MetricsTracker:
    """Tracks per‑step metrics and saves as JSON lines."""

    def __init__(self, filepath: str) -> None:
        self.filepath = filepath
        self.records: list[dict[str, Any]] = []
        os.makedirs(os.path.dirname(filepath), exist_ok=True)

    def log(self, metrics: dict[str, Any], step: int | None = None, flush: bool = False) -> None:
        """Log."""
        entry = {"step": step} if step is not None else {}
        entry.update(metrics)
        self.records.append(entry)
        if flush:
            self._flush()

    def _flush(self) -> None:
        """flush."""
        if self.filepath:
            with open(self.filepath, "a") as f:
                for rec in self.records:
                    f.write(json.dumps(rec) + "\n")
            self.records.clear()

    def save(self) -> None:
        """Save."""
        self._flush()

    def to_dataframe(self) -> pd.DataFrame:
        """To dataframe."""
        all_records = []
        if os.path.exists(self.filepath):
            with open(self.filepath) as f:
                for line in f:
                    try:
                        all_records.append(json.loads(line))
                    except json.JSONDecodeError:
                        pass
        all_records.extend(self.records)
        return pd.DataFrame(all_records)

Tracks per‑step metrics and saves as JSON lines.

Methods

def log(self,
metrics: dict[str, typing.Any],
step: int | None = None,
flush: bool = False) ‑> None

Expand source code

def log(self, metrics: dict[str, Any], step: int | None = None, flush: bool = False) -> None:
    """Log."""
    entry = {"step": step} if step is not None else {}
    entry.update(metrics)
    self.records.append(entry)
    if flush:
        self._flush()

Log.

def save(self) ‑> None

Expand source code

def save(self) -> None:
    """Save."""
    self._flush()

Save.

def to_dataframe(self) ‑> pandas.DataFrame

Expand source code

def to_dataframe(self) -> pd.DataFrame:
    """To dataframe."""
    all_records = []
    if os.path.exists(self.filepath):
        with open(self.filepath) as f:
            for line in f:
                try:
                    all_records.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    all_records.extend(self.records)
    return pd.DataFrame(all_records)

To dataframe.

class ModelComparison (results_dirs: list[str], output_dir: str = './results/comparison')

Expand source code

class ModelComparison:
    def __init__(self, results_dirs: list[str], output_dir: str = "./results/comparison") -> None:
        self.results_dirs = results_dirs
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.df_all = None

    def load_all_results(self) -> pd.DataFrame:
        """Load all results from disk or config."""
        dfs = []
        for d in self.results_dirs:
            csv_path = os.path.join(d, "performance.csv")
            if os.path.exists(csv_path):
                df = pd.read_csv(csv_path)
                df["run"] = os.path.basename(d)
                dfs.append(df)
        if dfs:
            self.df_all = pd.concat(dfs, ignore_index=True)
        return self.df_all

    def generate_all_plots(self) -> None:
        """Generate all plots."""
        if self.df_all is None:
            self.load_all_results()
        plots_dir = self.output_dir / "plots"
        create_comprehensive_plots(self.df_all, str(plots_dir))
        plot_tradeoff(self.df_all, str(plots_dir))
        print(f"Plots saved to {plots_dir}")

Methods

def generate_all_plots(self) ‑> None

Expand source code

def generate_all_plots(self) -> None:
    """Generate all plots."""
    if self.df_all is None:
        self.load_all_results()
    plots_dir = self.output_dir / "plots"
    create_comprehensive_plots(self.df_all, str(plots_dir))
    plot_tradeoff(self.df_all, str(plots_dir))
    print(f"Plots saved to {plots_dir}")

Generate all plots.

def load_all_results(self) ‑> pandas.DataFrame

Expand source code

def load_all_results(self) -> pd.DataFrame:
    """Load all results from disk or config."""
    dfs = []
    for d in self.results_dirs:
        csv_path = os.path.join(d, "performance.csv")
        if os.path.exists(csv_path):
            df = pd.read_csv(csv_path)
            df["run"] = os.path.basename(d)
            dfs.append(df)
    if dfs:
        self.df_all = pd.concat(dfs, ignore_index=True)
    return self.df_all

Load all results from disk or config.