analyze: add MoE router gating/share diagnostic plots
CI / Lint (ruff check) (push) Successful in 58s
CI / Format (ruff format) (push) Successful in 1m3s
CI / Type check (ty) (push) Successful in 1m10s
CI / Tests (push) Successful in 1m56s
CI / Bump version, build & publish wheel (push) Has been skipped

New "model" family in the gallery: router_gating (mean soft gate weight
vs. pre-step energy, showing the router's soft decision boundaries) and
router_share_by_pdg/router_share_by_process (stacked top-1 dispatch share
by species / true physics process). Needs a live checkpoint's Router, so
it's a documented exception to the rest of the package's polars/numpy-only
contract; gracefully degrades to a placeholder for non-MoE checkpoints.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-24 13:32:25 +02:00
parent 1115301eb9
commit 60c2ca1985
7 changed files with 592 additions and 7 deletions
+77
View File
@@ -137,12 +137,89 @@ def _render_bar(r: Reduced, params: dict):
return fig
def _render_router_gating(r: Reduced, params: dict):
n_experts = r.payload["n_experts"]
log_x = r.payload.get("log_x", False)
fig, axes = ps.new_figure(
"slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
)
flat = axes.ravel()
for ax, key in zip(flat, ("rollout", "reference")):
side = r.payload.get(key, {})
centers = np.asarray(side.get("centers", []))
means = np.asarray(side.get("means", []))
if len(centers) and means.size:
cum = np.zeros(len(centers))
for i in range(n_experts):
ax.fill_between(
centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
)
cum = cum + means[:, i]
if log_x:
ax.set_xscale("log")
ax.set_ylim(0, 1)
ax.set_title(_SERIES_LABELS[key], fontsize=8)
ax.set_xlabel(r.xlabel)
flat[0].set_ylabel("mean gate weight")
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
return fig
def _render_router_share(r: Reduced, params: dict):
categories = r.payload["categories"]
n_experts = r.payload["n_experts"]
x = np.arange(len(categories))
present = [k for k in ("rollout", "reference") if k in r.payload]
fig, axes = ps.new_figure(
"slide-16x9",
title=r.title,
params=params,
nrows=1,
ncols=len(present),
squeeze=False,
)
flat = axes.ravel()
for ax, key in zip(flat, present):
side = r.payload[key]
shares = np.array([side[c] for c in categories]) # (n_cat, n_experts)
bottom = np.zeros(len(categories))
for i in range(n_experts):
ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
bottom += shares[:, i]
ax.set_xticks(x)
ax.set_xticklabels(categories, rotation=45, ha="right")
ax.set_ylim(0, 1)
ax.set_title(_SERIES_LABELS[key], fontsize=8)
flat[0].set_ylabel("share of rows dispatched to expert")
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
return fig
def _render_unavailable(r: Reduced, params: dict):
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
ax.axis("off")
ax.text(
0.5,
0.5,
r.payload.get("note", "not available"),
ha="center",
va="center",
wrap=True,
fontsize=10,
transform=ax.transAxes,
)
return fig
_RENDERERS = {
"overlay_hist": _render_overlay,
"single_hist": _render_single,
"grouped_hist": _render_grouped,
"profile": _render_profile,
"bar": _render_bar,
"router_gating": _render_router_gating,
"router_share": _render_router_share,
"unavailable": _render_unavailable,
}