Add router gating diagnostic for MoE checkpoints

plot_router_gating visualizes soft expert gate weights vs. a continuous
routing axis (e.g. pre-step energy), binned into equal-population
quantiles and stacked to show the router's soft decision boundaries.
Wired into rollout_validation.ipynb as a new notebook-only section
that loads a checkpoint's Router directly, since gate weights aren't
present in rollout/predict parquet output.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/analysis.py

@@ -2294,3 +2294,66 @@ def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6):
         "length traveled share by particle type",
         max_slices,
     )
+
+
+def plot_router_gating(
+    x: np.ndarray,
+    gate_weights: np.ndarray,
+    x_label: str = "energy",
+    log_x: bool = True,
+    n_bins: int = 40,
+    figsize: tuple[float, float] = (7, 4),
+):
+    """Soft mixture-of-experts gate weight vs. a continuous routing axis.
+
+    Unlike everything else in this module, this doesn't stream from a
+    predict/rollout file — `gate_weights` (N, n_experts, rows already summing
+    to 1, `giant.model.network.Router.gate`'s contract) has to come from a
+    live `Router.gate(cond_cont, cond_cat)` call against a loaded checkpoint,
+    which is a deliberate exception to this module's file-only-diagnostics
+    design (see the module docstring); that on-the-fly step belongs in the
+    calling notebook, not here.
+
+    `x` is binned into `n_bins` equal-population (quantile) bins rather than
+    equal-width ones, since routing axes like energy are usually heavy-tailed
+    and equal-width bins would leave the upper end almost empty. One line per
+    expert, mean gate weight per bin stacked as filled areas — since rows of
+    `gate_weights` are a partition of unity, the stack always fills exactly
+    to 1, and the visible crossover bands are the router's soft decision
+    boundaries (where two experts' means cross ~0.5).
+    """
+    x_arr = np.asarray(x)
+    n_experts = gate_weights.shape[1]
+    order = np.argsort(x_arr)
+    x_sorted = x_arr[order]
+    gw_sorted = gate_weights[order]
+
+    edges = np.quantile(x_sorted, np.linspace(0, 1, n_bins + 1))
+    edges[-1] = np.nextafter(edges[-1], np.inf)  # include the max value
+    bin_idx = np.clip(np.digitize(x_sorted, edges[1:-1]), 0, n_bins - 1)
+
+    centers = np.full(n_bins, np.nan)
+    means = np.full((n_bins, n_experts), np.nan)
+    for b in range(n_bins):
+        mask = bin_idx == b
+        if mask.any():
+            centers[b] = x_sorted[mask].mean()
+            means[b] = gw_sorted[mask].mean(axis=0)
+
+    valid = ~np.isnan(centers)
+    centers, means = centers[valid], means[valid]
+
+    fig, ax = plt.subplots(figsize=figsize)
+    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_xlabel(x_label)
+    ax.set_ylabel("mean gate weight")
+    ax.set_ylim(0, 1)
+    ax.set_title("soft router gating")
+    ax.legend(fontsize=8, ncol=min(n_experts, 4))
+    fig.tight_layout()
+    return fig