Module 7. Mixture of Experts

import hashlib
import json
import math
from dataclasses import asdict, dataclass
from numbers import Integral, Real
from typing import Any, Literal

import importlib.util, subprocess, sys
if importlib.util.find_spec("numpy") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "numpy"])
import numpy as np

QWEN_CONFIG_JSON = '{\n    "architectures": [\n        "Qwen3_5MoeForConditionalGeneration"\n    ],\n    "image_token_id": 248056,\n    "model_type": "qwen3_5_moe",\n    "text_config": {\n        "attention_bias": false,\n        "attention_dropout": 0.0,\n        "attn_output_gate": true,\n        "dtype": "bfloat16",\n        "eos_token_id": 248044,\n        "full_attention_interval": 4,\n        "head_dim": 256,\n        "hidden_act": "silu",\n        "hidden_size": 2048,\n        "initializer_range": 0.02,\n        "layer_types": [\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention",\n            "linear_attention",\n            "linear_attention",\n            "linear_attention",\n            "full_attention"\n        ],\n        "linear_conv_kernel_dim": 4,\n        "linear_key_head_dim": 128,\n        "linear_num_key_heads": 16,\n        "linear_num_value_heads": 32,\n        "linear_value_head_dim": 128,\n        "max_position_embeddings": 262144,\n        "mlp_only_layers": [],\n        "model_type": "qwen3_5_moe_text",\n        "moe_intermediate_size": 512,\n        "mtp_num_hidden_layers": 1,\n        "mtp_use_dedicated_embeddings": false,\n        "num_attention_heads": 16,\n        "num_experts": 256,\n        "num_experts_per_tok": 8,\n        "num_hidden_layers": 40,\n        "num_key_value_heads": 2,\n        "rms_norm_eps": 1e-06,\n        "router_aux_loss_coef": 0.001,\n        "shared_expert_intermediate_size": 512,\n        "use_cache": true,\n        "vocab_size": 248320,\n        "mamba_ssm_dtype": "float32",\n        "rope_parameters": {\n            "mrope_interleaved": true,\n            "mrope_section": [\n                11,\n                11,\n                10\n            ],\n            "rope_type": "default",\n            "rope_theta": 10000000,\n            "partial_rotary_factor": 0.25\n        }\n    },\n    "tie_word_embeddings": false,\n    "transformers_version": "4.57.0.dev0",\n    "video_token_id": 248057,\n    "vision_config": {\n        "deepstack_visual_indexes": [],\n        "depth": 27,\n        "hidden_act": "gelu_pytorch_tanh",\n        "hidden_size": 1152,\n        "in_channels": 3,\n        "initializer_range": 0.02,\n        "intermediate_size": 4304,\n        "model_type": "qwen3_5_moe",\n        "num_heads": 16,\n        "num_position_embeddings": 2304,\n        "out_hidden_size": 2048,\n        "patch_size": 16,\n        "spatial_merge_size": 2,\n        "temporal_patch_size": 2\n    },\n    "vision_end_token_id": 248054,\n    "vision_start_token_id": 248053\n}'
QWEN_CONFIG_SHA256 = '5e4d7f74fec2f360eb9cfbfcd6ec0c4c76e684d3a11caaed259d9fd9bfbc7944'
assert hashlib.sha256(QWEN_CONFIG_JSON.encode("utf-8")).hexdigest() == QWEN_CONFIG_SHA256
QWEN_TEXT = json.loads(QWEN_CONFIG_JSON)["text_config"]
REFERENCE_JSON = '{"active_expert_law":{"E128_k4_N16":{"simulation":0.39857421875,"theory":0.3982896965679277},"E128_k4_N64":{"simulation":0.869296875,"theory":0.868915967521525},"E256_k8_N16":{"simulation":0.39908203125,"theory":0.3982896965679277},"E256_k8_N64":{"simulation":0.86845703125,"theory":0.868915967521525},"E896_k16_N16":{"simulation":0.2504827008928572,"theory":0.2504603703415571},"E896_k16_N64":{"simulation":0.6853069196428571,"theory":0.6843699099436731}},"balance":{"proportional":{"after":[0.031524658203125,0.032073974609375,0.0316162109375,0.03106689453125,0.0311279296875,0.0299072265625,0.0306396484375,0.030487060546875,0.0323486328125,0.031005859375,0.031280517578125,0.031463623046875,0.032012939453125,0.030548095703125,0.029998779296875,0.031982421875,0.031341552734375,0.030670166015625,0.031097412109375,0.031829833984375,0.03106689453125,0.03143310546875,0.032196044921875,0.03192138671875,0.03057861328125,0.03173828125,0.0313720703125,0.031005859375,0.030975341796875,0.03179931640625,0.030792236328125,0.031097412109375],"before":[0.0804443359375,0.08746337890625,0.0877685546875,0.079925537109375,0.0166015625,0.0250244140625,0.02374267578125,0.029022216796875,0.030731201171875,0.02301025390625,0.0234375,0.0252685546875,0.02923583984375,0.026336669921875,0.016143798828125,0.02020263671875,0.0311279296875,0.022613525390625,0.022796630859375,0.029205322265625,0.016387939453125,0.015655517578125,0.02044677734375,0.025482177734375,0.024139404296875,0.023834228515625,0.0242919921875,0.0242919921875,0.02203369140625,0.0283203125,0.02496337890625,0.020050048828125],"bias":[-0.03043212890625,-0.03636474609375004,-0.03786621093750004,-0.035046386718749996,0.021752929687499988,-0.0007324218750000033,0.003442382812499999,-0.004284667968750004,-0.002709960937500002,0.00487060546875,0.0048339843749999984,-0.0004028320312499989,-0.004650878906249997,-0.001684570312499999,0.017358398437499988,0.013403320312500003,-0.006701660156249992,0.0047241210937499995,0.006225585937499997,-0.004980468749999995,0.016809082031250006,0.024023437500000022,0.013293457031249997,0.0034057617187500007,0.0036987304687499984,0.0024902343750000003,0.0012451171875000001,0.0020141601562500035,0.00974121093750001,-0.0007324218749999982,0.002380371093750004,0.01087646484375],"cv_after":0.01932878458289989,"cv_before":0.6505852027323183,"mode":"proportional","rate":0.15,"uniform":0.03125},"sign":{"after":[0.03179931640625,0.035308837890625,0.031646728515625,0.02972412109375,0.030487060546875,0.02886962890625,0.03021240234375,0.031341552734375,0.0328369140625,0.03057861328125,0.0318603515625,0.03082275390625,0.032928466796875,0.030517578125,0.029632568359375,0.033050537109375,0.03094482421875,0.030364990234375,0.030914306640625,0.032379150390625,0.0318603515625,0.0316162109375,0.031219482421875,0.030731201171875,0.029998779296875,0.02984619140625,0.031585693359375,0.030853271484375,0.03106689453125,0.032684326171875,0.031219482421875,0.031097412109375],"before":[0.0804443359375,0.08746337890625,0.0877685546875,0.079925537109375,0.0166015625,0.0250244140625,0.02374267578125,0.029022216796875,0.030731201171875,0.02301025390625,0.0234375,0.0252685546875,0.02923583984375,0.026336669921875,0.016143798828125,0.02020263671875,0.0311279296875,0.022613525390625,0.022796630859375,0.029205322265625,0.016387939453125,0.015655517578125,0.02044677734375,0.025482177734375,0.024139404296875,0.023834228515625,0.0242919921875,0.0242919921875,0.02203369140625,0.0283203125,0.02496337890625,0.020050048828125],"bias":[-0.014000000000000005,-0.01800000000000001,-0.02100000000000001,-0.01900000000000001,0.037000000000000026,0.014000000000000005,0.01900000000000001,0.013000000000000005,0.014000000000000005,0.02000000000000001,0.022000000000000013,0.015000000000000006,0.013000000000000005,0.015000000000000006,0.03300000000000002,0.03200000000000002,0.009000000000000001,0.02100000000000001,0.022000000000000013,0.012000000000000004,0.035000000000000024,0.04100000000000003,0.028000000000000018,0.01800000000000001,0.01900000000000001,0.016000000000000007,0.01800000000000001,0.01800000000000001,0.026000000000000016,0.017000000000000008,0.01900000000000001,0.027000000000000017],"cv_after":0.03931118956596903,"cv_before":0.6505852027323183,"mode":"sign","rate":0.001,"uniform":0.03125}},"capacity":{"after_counts":[6,4,4,4],"before_counts":[12,4,4,4],"capacity_per_expert":6,"dropped_assignments":6,"fully_dropped_tokens":[],"logits":[[4.0,3.0,0.0,-1.0],[4.0,0.0,3.0,-1.0],[4.0,0.0,-1.0,3.0],[4.0,3.0,0.0,-1.0],[4.0,0.0,3.0,-1.0],[4.0,0.0,-1.0,3.0],[4.0,3.0,0.0,-1.0],[4.0,0.0,3.0,-1.0],[4.0,0.0,-1.0,3.0],[4.0,3.0,0.0,-1.0],[4.0,0.0,3.0,-1.0],[4.0,0.0,-1.0,3.0]],"policy":"fifo_drop_then_renormalize_survivors"},"combinations":{"C128_4":10668000,"C256_8_unconstrained_upper_bound":409663695276000,"C8_2":28},"expert_parallel":{"bytes_per_value":2,"hidden_size":7168,"round_trip_bytes":469762048,"tokens":2048,"top_k":8},"latent_moe":{"kimi_k3_hidden":7168,"kimi_k3_latent":3584,"payload_reduction":2.0,"routed_experts":896,"selected_experts":16,"shared_experts":2},"losses":{"aux_logits":[[4.0,1.0,0.0,-1.0],[3.0,1.0,0.0,-1.0],[4.0,0.0,1.0,-1.0],[3.0,0.0,1.0,-1.0]],"switch_aux":3.522646252388999,"z_loss":13.363987745409077},"models":{"deepseek_v3_hf_package_total":685000000000,"deepseek_v3_official_active":37000000000,"deepseek_v3_official_main_total":671000000000,"deepseek_v3_reconstruction":{"active_parameters":36625618432,"attention":11413547008,"dense_ffn":1189085184,"extra_active":0,"extra_total":0,"norms":881664,"routed_and_shared_experts_active":22988980224,"routed_and_shared_experts_total":656463101952,"router":106445312,"sinks":0,"total_parameters":671026419200,"vocab_active":926679040,"vocab_total":1853358080},"gqa_params_per_layer_with_bias":26550080,"mla_params_per_layer":187107328,"oss120":{"active_parameters":5132849472,"attention":955802880,"dense_ffn":0,"extra_active":0,"extra_total":0,"norms":210240,"routed_and_shared_experts_active":3584424960,"routed_and_shared_experts_total":114701598720,"router":13275648,"sinks":2304,"total_parameters":116829156672,"vocab_active":579133440,"vocab_total":1158266880},"oss20":{"active_parameters":3608307264,"attention":637201920,"dense_ffn":0,"extra_active":0,"extra_total":0,"norms":141120,"routed_and_shared_experts_active":2389616640,"routed_and_shared_experts_total":19116933120,"router":2212608,"sinks":1536,"total_parameters":20914757184,"vocab_active":579133440,"vocab_total":1158266880}},"qwen35_config":{"hidden_size":2048,"layers":40,"mtp_layers":1,"routed_experts":256,"routed_intermediate":512,"selected_experts":8,"shared_experts":1,"shared_intermediate":512},"router":{"bias":[0.0,0.0,0.0,0.35,0.0,0.0,0.0,0.0],"boundary_update":{"one_step":{"e0":0.7497747834990175,"e6":0.7296995026250039},"two_steps":{"e0":0.7310247834990176,"e6":0.7484495026250039}},"gates":[0.2694994612569363,0.25980636470427615,0.24657595902244958,0.224118215016338],"gates_with_bias":[0.18368780089939873,0.28354280527375636,0.27334461127547466,0.2594247825513703],"logits":[1.2,-0.3,2.1,0.4,1.7,-1.0,0.9,2.5],"scores":[0.7685247834990175,0.425557483188341,0.8909031788043871,0.598687660112452,0.8455347349164652,0.2689414213699951,0.7109495026250039,0.9241418199787566],"selected":[7,2,4,0],"selected_with_bias":[3,7,2,4]}}'
REFERENCE_SHA256 = 'de9ba39bbe902907f8bbae0e6bfe1a4d19028eeedc2f167e026995a4f85ddeed'
REFERENCE = json.loads(REFERENCE_JSON)
assert hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest() == REFERENCE_SHA256
print('contract verified', REFERENCE_SHA256[:16])


def _finite_array(name: str, value: Any, *, ndim: int | None = None) -> np.ndarray:
    """Return ``value`` as a finite float64 array with an optional rank contract."""

    array = np.asarray(value, dtype=np.float64)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}; got shape {array.shape}")
    if array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def _positive_int(name: str, value: Any, *, minimum: int = 1) -> int:
    """Validate an integer scalar while rejecting booleans explicitly."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result < minimum:
        raise ValueError(f"{name} must be >= {minimum}")
    return result


def _bounded_real(
    name: str,
    value: Any,
    *,
    lower: float,
    upper: float,
    lower_inclusive: bool = True,
    upper_inclusive: bool = True,
) -> float:
    """Validate one finite real scalar against an explicit interval."""

    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"{name} must be finite")
    lower_ok = result >= lower if lower_inclusive else result > lower
    upper_ok = result <= upper if upper_inclusive else result < upper
    if not (lower_ok and upper_ok):
        left = "[" if lower_inclusive else "("
        right = "]" if upper_inclusive else ")"
        raise ValueError(f"{name} must lie in {left}{lower}, {upper}{right}")
    return result


def stable_sigmoid(value: Any) -> np.ndarray:
    """Evaluate the logistic sigmoid without overflowing for large finite logits."""

    x = _finite_array("value", value)
    out = np.empty_like(x)
    nonnegative = x >= 0.0
    out[nonnegative] = 1.0 / (1.0 + np.exp(-x[nonnegative]))
    exp_x = np.exp(x[~nonnegative])
    out[~nonnegative] = exp_x / (1.0 + exp_x)
    return out


def stable_softmax(value: Any, *, axis: int = -1) -> np.ndarray:
    """Compute a finite softmax along ``axis`` with max subtraction."""

    x = _finite_array("value", value)
    if axis < -x.ndim or axis >= x.ndim:
        raise ValueError("axis is outside the array rank")
    shifted = x - np.max(x, axis=axis, keepdims=True)
    exp_shifted = np.exp(shifted)
    denominator = exp_shifted.sum(axis=axis, keepdims=True)
    if np.any(denominator <= 0.0) or not np.all(np.isfinite(denominator)):
        raise FloatingPointError("softmax normalization failed")
    return exp_shifted / denominator


def _log_sigmoid(value: np.ndarray) -> np.ndarray:
    """Return log(sigmoid(value)) in a numerically stable form."""

    return -np.logaddexp(0.0, -value)


def _deterministic_top_k(scores: np.ndarray, top_k: int) -> np.ndarray:
    """Select descending scores with lower expert index breaking exact ties."""

    token_count, expert_count = scores.shape
    indices = np.arange(expert_count)
    selected = np.empty((token_count, top_k), dtype=np.int64)
    for token in range(token_count):
        order = np.lexsort((indices, -scores[token]))
        selected[token] = order[:top_k]
    return selected


@dataclass(frozen=True)
class RouterResult:
    """Validated output of one token-choice router pass.

    Shapes are ``clean_scores:[T,E]``, ``selected_experts:[T,k]``, and
    ``gate_weights:[T,k]``. ``selection_scores`` includes the optional balancing bias;
    ``gate_weights`` never does.
    """

    clean_scores: np.ndarray
    selection_scores: np.ndarray
    selected_experts: np.ndarray
    gate_weights: np.ndarray
    expert_counts: np.ndarray


@dataclass(frozen=True)
class DispatchPlan:
    """Capacity-limited token-to-expert assignments under a deterministic FIFO policy."""

    capacity_per_expert: int
    accepted: np.ndarray
    gate_weights: np.ndarray
    expert_counts: np.ndarray
    dropped_assignments: int
    fully_dropped_tokens: tuple[int, ...]


@dataclass(frozen=True)
class LoadStatistics:
    """Fractions and imbalance diagnostics for expert assignment counts."""

    fractions: np.ndarray
    coefficient_of_variation: float
    maximum_fraction: float
    minimum_fraction: float

def route_logits(
    logits: Any,
    *,
    top_k: int,
    scoring: Literal["sigmoid", "softmax"] = "sigmoid",
    selection_bias: Any | None = None,
) -> RouterResult:
    """Route a non-empty token batch with deterministic token-choice top-k.

    Parameters
    ----------
    logits:
        Router logits of shape ``[T,E]`` or one token of shape ``[E]``.
    top_k:
        Number of routed experts per token, constrained to ``1 <= k <= E``.
    scoring:
        ``"sigmoid"`` for independent V3-style scores or ``"softmax"`` for a
        categorical router.
    selection_bias:
        Optional vector ``[E]`` added **only** to the selection scores. Gate weights are
        normalized from the clean scores, preserving the aux-loss-free contract.

    Exact ties are broken by lower expert index so CPU, Colab, and different NumPy builds
    produce the same route.
    """

    raw = _finite_array("logits", logits)
    if raw.ndim == 1:
        raw = raw[None, :]
    if raw.ndim != 2:
        raise ValueError("logits must have shape [T,E] or [E]")
    token_count, expert_count = raw.shape
    k = _positive_int("top_k", top_k)
    if k > expert_count:
        raise ValueError("top_k cannot exceed the number of experts")
    if scoring not in {"sigmoid", "softmax"}:
        raise ValueError("scoring must be 'sigmoid' or 'softmax'")

    if selection_bias is None:
        bias = np.zeros(expert_count, dtype=np.float64)
    else:
        bias = _finite_array("selection_bias", selection_bias, ndim=1)
        if bias.shape != (expert_count,):
            raise ValueError("selection_bias must have shape [E]")

    clean_scores = stable_sigmoid(raw) if scoring == "sigmoid" else stable_softmax(raw, axis=1)
    selection_scores = clean_scores + bias[None, :]
    # Both score functions preserve rank mathematically. Ranking raw logits when
    # no bias is supplied avoids artificial ties from saturation/underflow.
    selected = _deterministic_top_k(raw if selection_bias is None else selection_scores, k)

    selected_logits = np.take_along_axis(raw, selected, axis=1)
    if scoring == "sigmoid":
        # Normalize selected sigmoid scores in log space. This remains finite even when
        # every displayed sigmoid underflows to zero at extreme negative logits.
        log_weights = _log_sigmoid(selected_logits)
        gate_weights = stable_softmax(log_weights, axis=1)
    else:
        # Re-normalizing clean softmax scores over top-k is exactly softmax of selected logits.
        gate_weights = stable_softmax(selected_logits, axis=1)

    counts = np.bincount(selected.ravel(), minlength=expert_count).astype(np.int64)
    if not np.allclose(gate_weights.sum(axis=1), 1.0, atol=1e-12):
        raise AssertionError("router gates must sum to one per token")
    return RouterResult(
        clean_scores=clean_scores,
        selection_scores=selection_scores,
        selected_experts=selected,
        gate_weights=gate_weights,
        expert_counts=counts,
    )


def expert_capacity(
    *,
    token_count: int,
    top_k: int,
    expert_count: int,
    capacity_factor: Real,
) -> int:
    """Return ``ceil(capacity_factor * T * k / E)`` assignments per expert."""

    tokens = _positive_int("token_count", token_count)
    k = _positive_int("top_k", top_k)
    experts = _positive_int("expert_count", expert_count)
    if k > experts:
        raise ValueError("top_k cannot exceed expert_count")
    factor = _bounded_real(
        "capacity_factor",
        capacity_factor,
        lower=0.0,
        upper=float("inf"),
        lower_inclusive=False,
    )
    return max(1, math.ceil(factor * tokens * k / experts))


def dispatch_with_capacity(
    router: RouterResult,
    *,
    capacity_factor: Real = 1.0,
    renormalize_kept_gates: bool = True,
) -> DispatchPlan:
    """Apply deterministic FIFO expert capacity and optionally renormalize kept gates.

    This transparent policy is intentionally simpler than production grouped dispatch. Each
    token is visited in input order and each route slot in router order. Overflow assignments
    are dropped; no hidden rerouting occurs. If ``renormalize_kept_gates`` is true, surviving
    weights are renormalized as an explicit *course policy* rather than a universal MoE rule.
    Tokens with no surviving assignment emit zero from the expert branch, remain on the
    residual path outside this educational layer, and are reported explicitly.
    """

    selected = np.asarray(router.selected_experts, dtype=np.int64)
    gates = _finite_array("router.gate_weights", router.gate_weights, ndim=2)
    if selected.ndim != 2 or selected.shape != gates.shape:
        raise ValueError("router selected_experts and gate_weights must share shape [T,k]")
    token_count, k = selected.shape
    if token_count == 0:
        raise ValueError("router output must contain at least one token")
    expert_count = int(np.asarray(router.expert_counts).size)
    if expert_count < 1 or np.any(selected < 0) or np.any(selected >= expert_count):
        raise ValueError("router contains an invalid expert index")

    capacity = expert_capacity(
        token_count=token_count,
        top_k=k,
        expert_count=expert_count,
        capacity_factor=capacity_factor,
    )
    accepted = np.zeros_like(selected, dtype=bool)
    counts = np.zeros(expert_count, dtype=np.int64)
    for token in range(token_count):
        for slot in range(k):
            expert = int(selected[token, slot])
            if counts[expert] < capacity:
                accepted[token, slot] = True
                counts[expert] += 1

    effective_gates = np.where(accepted, gates, 0.0)
    kept_mass = effective_gates.sum(axis=1, keepdims=True)
    fully_dropped = tuple(np.flatnonzero(kept_mass[:, 0] == 0.0).tolist())
    if renormalize_kept_gates:
        np.divide(
            effective_gates,
            kept_mass,
            out=effective_gates,
            where=kept_mass > 0.0,
        )
    dropped = int(accepted.size - accepted.sum())
    return DispatchPlan(
        capacity_per_expert=capacity,
        accepted=accepted,
        gate_weights=effective_gates,
        expert_counts=counts,
        dropped_assignments=dropped,
        fully_dropped_tokens=fully_dropped,
    )


def load_statistics(counts: Any) -> LoadStatistics:
    """Convert non-negative assignment counts into fractions and load CV."""

    raw = _finite_array("counts", counts, ndim=1)
    if np.any(raw < 0.0):
        raise ValueError("counts must be non-negative")
    total = float(raw.sum())
    if total <= 0.0:
        raise ValueError("at least one assignment is required")
    fractions = raw / total
    mean = float(fractions.mean())
    return LoadStatistics(
        fractions=fractions,
        coefficient_of_variation=float(fractions.std() / mean),
        maximum_fraction=float(fractions.max()),
        minimum_fraction=float(fractions.min()),
    )


def aux_free_bias_step(
    selection_bias: Any,
    counts: Any,
    *,
    rate: Real,
    mode: Literal["sign", "proportional"] = "sign",
) -> np.ndarray:
    """Update expert selection biases outside backpropagation.

    ``mode="sign"`` reproduces the main update rule in Wang et al. (2024): an
    overloaded expert moves down by one fixed step and an underloaded expert moves up.
    ``mode="proportional"`` is the smoother ablation/teaching controller used in this
    module's convergence trace. Both modify selection only; clean gate values remain
    unchanged.
    """

    bias = _finite_array("selection_bias", selection_bias, ndim=1)
    stats = load_statistics(counts)
    if stats.fractions.shape != bias.shape:
        raise ValueError("selection_bias and counts must have the same shape")
    step = _bounded_real("rate", rate, lower=0.0, upper=float("inf"), lower_inclusive=False)
    error = stats.fractions.mean() - stats.fractions
    if mode == "sign":
        delta = np.sign(error)
    elif mode == "proportional":
        delta = error
    else:
        raise ValueError("mode must be 'sign' or 'proportional'")
    return bias + step * delta


def switch_auxiliary_loss(logits: Any, *, top_k: int = 1) -> float:
    """Compute ``E * sum_i f_i P_i`` for a softmax token-choice router batch."""

    raw = _finite_array("logits", logits)
    if raw.ndim != 2:
        raise ValueError("logits must have shape [T,E]")
    routed = route_logits(raw, top_k=top_k, scoring="softmax")
    token_count, expert_count = raw.shape
    assignment_fraction = routed.expert_counts / (token_count * top_k)
    mean_probability = routed.clean_scores.mean(axis=0)
    return float(expert_count * np.dot(assignment_fraction, mean_probability))


def router_z_loss(logits: Any) -> float:
    """Compute the ST-MoE router z-loss ``mean(logsumexp(logits)^2)``."""

    raw = _finite_array("logits", logits)
    if raw.ndim != 2:
        raise ValueError("logits must have shape [T,E]")
    maximum = np.max(raw, axis=1, keepdims=True)
    log_partition = maximum[:, 0] + np.log(np.exp(raw - maximum).sum(axis=1))
    return float(np.mean(log_partition**2))


ref = REFERENCE["router"]
base = route_logits(np.array(ref["logits"]), top_k=4, scoring="sigmoid")
biased = route_logits(np.array(ref["logits"]), top_k=4, scoring="sigmoid", selection_bias=ref["bias"])
assert base.selected_experts[0].tolist() == ref["selected"]
assert np.allclose(base.gate_weights[0], ref["gates"], atol=1e-12)
assert biased.selected_experts[0].tolist() == ref["selected_with_bias"]
assert np.allclose(biased.gate_weights[0], ref["gates_with_bias"], atol=1e-12)
assert route_logits(np.zeros(4), top_k=2, scoring="softmax").selected_experts[0].tolist() == [0, 1]
extreme = route_logits([-1002., -1000., -1001.], top_k=2)
assert extreme.selected_experts.tolist() == [[1, 2]]
assert np.all(np.isfinite(extreme.gate_weights))
assert np.allclose(extreme.gate_weights, [[1/(1+np.exp(-1)), 1/(1+np.exp(1))]])
print("top-k", base.selected_experts[0].tolist(), "biased", biased.selected_experts[0].tolist())
print('router verified')

ref = REFERENCE["capacity"]
routed = route_logits(np.array(ref["logits"]), top_k=2, scoring="softmax")
plan = dispatch_with_capacity(routed, capacity_factor=1.0, renormalize_kept_gates=True)
assert ref["policy"] == "fifo_drop_then_renormalize_survivors"
assert plan.capacity_per_expert == ref["capacity_per_expert"] == 6
assert routed.expert_counts.tolist() == ref["before_counts"]
assert plan.expert_counts.tolist() == ref["after_counts"]
assert plan.dropped_assignments == ref["dropped_assignments"] == 6
assert list(plan.fully_dropped_tokens) == ref["fully_dropped_tokens"]
assert np.allclose(plan.gate_weights.sum(axis=1), 1.0)
print("before", routed.expert_counts, "after", plan.expert_counts, "dropped", plan.dropped_assignments)
# Fully dropped routed tokens remain observable; no division by zero.
all_same = route_logits(np.tile([9., 0.], (4, 1)), top_k=1)
limited = dispatch_with_capacity(all_same, capacity_factor=0.5)
assert limited.fully_dropped_tokens == (1, 2, 3)
assert np.all(limited.gate_weights[1:] == 0)
assert limited.accepted[:, 0].tolist() == [True, False, False, False]
print('capacity policy verified')

def swiglu_expert(inputs: Any, gate: Any, up: Any, down: Any) -> np.ndarray:
    """Apply one dense SwiGLU expert to ``inputs:[T,h]`` with explicit shape checks."""

    x = _finite_array("inputs", inputs)
    if x.ndim == 1:
        x = x[None, :]
    if x.ndim != 2:
        raise ValueError("inputs must have shape [T,h] or [h]")
    w_gate = _finite_array("gate", gate, ndim=2)
    w_up = _finite_array("up", up, ndim=2)
    w_down = _finite_array("down", down, ndim=2)
    hidden = x.shape[1]
    if w_gate.shape[0] != hidden or w_up.shape != w_gate.shape:
        raise ValueError("gate and up must both have shape [h,i_exp]")
    if w_down.shape != (w_gate.shape[1], hidden):
        raise ValueError("down must have shape [i_exp,h]")
    gate_values = x @ w_gate
    up_values = x @ w_up
    return (gate_values * stable_sigmoid(gate_values) * up_values) @ w_down


@dataclass(frozen=True)
class MoEForwardResult:
    """Output plus routing metadata from the transparent educational MoE layer."""

    output: np.ndarray
    router: RouterResult
    dispatch: DispatchPlan | None


class EducationalMoE:
    """Small NumPy MoE with routed and always-on shared SwiGLU experts.

    Expert tensors use shape ``[E,h,i_exp]`` for gate/up and ``[E,i_exp,h]`` for
    down. The Python expert loop is intentionally inspectable; production code would use
    grouped GEMMs, fused dispatch, and distributed expert parallelism.
    """

    def __init__(
        self,
        router_weights: Any,
        expert_gate: Any,
        expert_up: Any,
        expert_down: Any,
        *,
        shared_gate: Any | None = None,
        shared_up: Any | None = None,
        shared_down: Any | None = None,
    ) -> None:
        """Validate and store routed/shared expert tensors without hidden casting later."""

        self.router_weights = _finite_array("router_weights", router_weights, ndim=2)
        self.expert_gate = _finite_array("expert_gate", expert_gate, ndim=3)
        self.expert_up = _finite_array("expert_up", expert_up, ndim=3)
        self.expert_down = _finite_array("expert_down", expert_down, ndim=3)
        experts, hidden, intermediate = self.expert_gate.shape
        if self.router_weights.shape != (hidden, experts):
            raise ValueError("router_weights must have shape [h,E]")
        if self.expert_up.shape != self.expert_gate.shape:
            raise ValueError("expert_up must match expert_gate")
        if self.expert_down.shape != (experts, intermediate, hidden):
            raise ValueError("expert_down must have shape [E,i_exp,h]")

        supplied = (shared_gate is not None, shared_up is not None, shared_down is not None)
        if len(set(supplied)) != 1:
            raise ValueError("shared gate/up/down must be supplied together")
        if shared_gate is None:
            self.shared_gate = np.empty((0, hidden, intermediate), dtype=np.float64)
            self.shared_up = np.empty((0, hidden, intermediate), dtype=np.float64)
            self.shared_down = np.empty((0, intermediate, hidden), dtype=np.float64)
        else:
            self.shared_gate = _finite_array("shared_gate", shared_gate, ndim=3)
            self.shared_up = _finite_array("shared_up", shared_up, ndim=3)
            self.shared_down = _finite_array("shared_down", shared_down, ndim=3)
            shared = self.shared_gate.shape[0]
            if self.shared_gate.shape[1:] != (hidden, intermediate):
                raise ValueError("shared_gate must have shape [S,h,i_exp]")
            if self.shared_up.shape != self.shared_gate.shape:
                raise ValueError("shared_up must match shared_gate")
            if self.shared_down.shape != (shared, intermediate, hidden):
                raise ValueError("shared_down must have shape [S,i_exp,h]")

    def forward(
        self,
        inputs: Any,
        *,
        top_k: int,
        scoring: Literal["sigmoid", "softmax"] = "sigmoid",
        selection_bias: Any | None = None,
        capacity_factor: Real | None = None,
    ) -> MoEForwardResult:
        """Route and combine one batch; capacity overflow follows ``dispatch_with_capacity``."""

        x = _finite_array("inputs", inputs)
        if x.ndim == 1:
            x = x[None, :]
        if x.ndim != 2 or x.shape[1] != self.router_weights.shape[0]:
            raise ValueError("inputs must have shape [T,h]")
        router = route_logits(
            x @ self.router_weights,
            top_k=top_k,
            scoring=scoring,
            selection_bias=selection_bias,
        )
        dispatch = (
            dispatch_with_capacity(router, capacity_factor=capacity_factor)
            if capacity_factor is not None
            else None
        )
        gates = dispatch.gate_weights if dispatch is not None else router.gate_weights
        accepted = dispatch.accepted if dispatch is not None else np.ones_like(gates, dtype=bool)
        output = np.zeros_like(x)

        for expert in range(self.expert_gate.shape[0]):
            token_rows, slots = np.nonzero((router.selected_experts == expert) & accepted)
            if token_rows.size:
                expert_output = swiglu_expert(
                    x[token_rows],
                    self.expert_gate[expert],
                    self.expert_up[expert],
                    self.expert_down[expert],
                )
                output[token_rows] += gates[token_rows, slots, None] * expert_output

        for shared in range(self.shared_gate.shape[0]):
            output += swiglu_expert(
                x,
                self.shared_gate[shared],
                self.shared_up[shared],
                self.shared_down[shared],
            )
        return MoEForwardResult(output=output, router=router, dispatch=dispatch)


rng = np.random.default_rng(60)
hidden, intermediate = 8, 12
x = rng.standard_normal((16, hidden))
router_w = rng.standard_normal((hidden, 1)) / math.sqrt(hidden)
gate = rng.standard_normal((1, hidden, intermediate)) / math.sqrt(hidden)
up = rng.standard_normal((1, hidden, intermediate)) / math.sqrt(hidden)
down = rng.standard_normal((1, intermediate, hidden)) / math.sqrt(intermediate)
layer = EducationalMoE(router_w, gate, up, down)
result = layer.forward(x, top_k=1)
dense = swiglu_expert(x, gate[0], up[0], down[0])
assert np.allclose(result.output, dense, atol=1e-12)
sg = rng.standard_normal((1, hidden, intermediate)) / math.sqrt(hidden)
su = rng.standard_normal((1, hidden, intermediate)) / math.sqrt(hidden)
sd = rng.standard_normal((1, intermediate, hidden)) / math.sqrt(intermediate)
with_shared = EducationalMoE(router_w, gate, up, down, shared_gate=sg, shared_up=su, shared_down=sd)
assert np.allclose(with_shared.forward(x, top_k=1).output, dense + swiglu_expert(x, sg[0], su[0], sd[0]))
print('MoE invariants verified')

def simulate_aux_free_balancing(*, mode: Literal["sign", "proportional"]) -> dict[str, Any]:
    """Reproduce deterministic 32-expert balancing traces.

    ``sign`` mirrors the main paper rule with a fixed-size correction. ``proportional`` is the
    smoother teaching controller used to visualize convergence; the paper also studies a
    proportional variant in an expert-parallel setting.
    """

    rng = np.random.default_rng(6)
    experts, top_k, token_count = 32, 4, 8192
    weights = rng.standard_normal((64, experts)) * 0.35
    constant = np.zeros(experts)
    constant[:4] = 2.5
    inputs = rng.standard_normal((token_count, 64))
    all_scores = stable_sigmoid(inputs @ weights + constant)

    def loads(bias: np.ndarray, scores: np.ndarray) -> np.ndarray:
        selected = _deterministic_top_k(scores + bias[None, :], top_k)
        return np.bincount(selected.ravel(), minlength=experts) / (scores.shape[0] * top_k)

    bias = np.zeros(experts)
    before = loads(bias, all_scores)
    rate = 0.001 if mode == "sign" else 0.15
    for _ in range(400):
        sample = all_scores[rng.integers(0, token_count, size=1024)]
        fractions = loads(bias, sample)
        bias = aux_free_bias_step(bias, fractions, rate=rate, mode=mode)
    after = loads(bias, all_scores)
    return {
        "mode": mode,
        "rate": rate,
        "before": before.tolist(),
        "after": after.tolist(),
        "bias": bias.tolist(),
        "uniform": 1.0 / experts,
        "cv_before": float(before.std() / before.mean()),
        "cv_after": float(after.std() / after.mean()),
    }


loss_ref = REFERENCE["losses"]
aux_logits = np.array(loss_ref["aux_logits"])
assert math.isclose(switch_auxiliary_loss(aux_logits), loss_ref["switch_aux"], rel_tol=1e-12)
assert math.isclose(router_z_loss(aux_logits), loss_ref["z_loss"], rel_tol=1e-12)
# Square before the token mean, not after it.
z = np.array([[0., 0.], [2., 2.]])
independent = (math.log(2)**2 + (2+math.log(2))**2)/2
assert math.isclose(router_z_loss(z), independent, rel_tol=1e-12)
assert not math.isclose(router_z_loss(z), (1+math.log(2))**2)
# The main paper rule is a fixed sign step.
bias = aux_free_bias_step(np.zeros(4), [8., 0., 0., 0.], rate=0.1, mode="sign")
assert np.allclose(bias, [-0.1, 0.1, 0.1, 0.1])
for mode in ("sign", "proportional"):
    got = simulate_aux_free_balancing(mode=mode)
    ref = REFERENCE["balance"][mode]
    assert np.allclose(got["before"], ref["before"], atol=1e-12)
    assert np.allclose(got["after"], ref["after"], atol=1e-12)
    assert math.isclose(got["cv_after"], ref["cv_after"], rel_tol=1e-12)
    print(mode, f"CV {got['cv_before']:.3f} -> {got['cv_after']:.3f}")
print('balancing verified')

@dataclass(frozen=True)
class MoEModelSpec:
    """Explicit parameter-counting contract for a dense/MoE Transformer.

    The class deliberately avoids pretending that one formula covers every MoE family. The
    caller states how many dense and MoE layers exist, whether the vocabulary matrices are tied,
    whether expert/router biases are present, and how many *full* vocabulary matrices are counted
    in the active-parameter convention. ``active_vocab_matrices=1`` is the compute-oriented
    convention used for gpt-oss: the dense output projection is counted, while an indexed input
    embedding lookup is not charged as a full matrix.
    """

    vocab_size: int
    hidden_size: int
    layers: int
    attention_params_per_layer: int
    routed_experts: int
    experts_per_token: int
    expert_intermediate: int
    shared_experts: int = 0
    dense_layers: int = 0
    dense_intermediate: int = 0
    vocab_matrices_total: int = 2
    active_vocab_matrices: int = 1
    router_bias: bool = False
    expert_bias: bool = False
    dense_bias: bool = False
    sink_params_per_layer: int = 0
    extra_total_parameters: int = 0
    extra_active_parameters: int = 0

    def __post_init__(self) -> None:
        for name in (
            "vocab_size", "hidden_size", "layers", "attention_params_per_layer",
            "routed_experts", "experts_per_token", "expert_intermediate",
            "vocab_matrices_total",
        ):
            _positive_int(name, getattr(self, name))
        for name in (
            "shared_experts", "dense_layers", "dense_intermediate",
            "active_vocab_matrices", "sink_params_per_layer",
            "extra_total_parameters", "extra_active_parameters",
        ):
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, Integral) or int(value) < 0:
                raise ValueError(f"{name} must be a non-negative integer")
        if self.experts_per_token > self.routed_experts:
            raise ValueError("experts_per_token cannot exceed routed_experts")
        if self.dense_layers > self.layers:
            raise ValueError("dense_layers cannot exceed layers")
        if self.dense_layers and self.dense_intermediate == 0:
            raise ValueError("dense_intermediate is required when dense_layers > 0")
        if self.active_vocab_matrices > self.vocab_matrices_total:
            raise ValueError("active_vocab_matrices cannot exceed vocab_matrices_total")


@dataclass(frozen=True)
class ParameterBreakdown:
    """Total and active counts under one declared convention."""

    total_parameters: int
    active_parameters: int
    vocab_total: int
    vocab_active: int
    attention: int
    routed_and_shared_experts_total: int
    routed_and_shared_experts_active: int
    dense_ffn: int
    router: int
    norms: int
    sinks: int
    extra_total: int
    extra_active: int

    @property
    def total_to_active_ratio(self) -> float:
        return self.total_parameters / self.active_parameters


def gqa_attention_parameters(
    *,
    hidden_size: int,
    query_heads: int,
    kv_heads: int,
    head_width: int,
    include_bias: bool = False,
) -> int:
    """Count Q/K/V/O projections for one GQA layer.

    With ``include_bias=True`` the count matches the fused QKV bias and output bias in the
    public gpt-oss reference implementation.
    """

    hidden = _positive_int("hidden_size", hidden_size)
    q_heads = _positive_int("query_heads", query_heads)
    k_heads = _positive_int("kv_heads", kv_heads)
    width = _positive_int("head_width", head_width)
    if k_heads > q_heads:
        raise ValueError("kv_heads cannot exceed query_heads")
    q_dim = q_heads * width
    kv_dim = k_heads * width
    weights = hidden * q_dim + 2 * hidden * kv_dim + q_dim * hidden
    biases = (q_dim + 2 * kv_dim + hidden) if include_bias else 0
    return int(weights + biases)


def deepseek_v3_mla_attention_parameters() -> int:
    """Course reconstruction of one DeepSeek-V3 MLA layer from public configuration fields.

    Projection biases are absent in the published config. The final two terms are the Q/latent
    normalization scales. This is an explicit reconstruction, not a claim that the public 671B
    headline follows a unique counting convention.
    """

    hidden, q_lora, kv_lora, rope, heads, head_width = 7168, 1536, 512, 64, 128, 128
    return int(
        hidden * q_lora
        + q_lora * (heads * (head_width + rope))
        + hidden * (kv_lora + rope)
        + kv_lora * (heads * (head_width + head_width))
        + heads * head_width * hidden
        + q_lora
        + kv_lora
    )


def expert_parameters(*, hidden_size: int, intermediate_size: int, include_bias: bool) -> int:
    """Count one SwiGLU expert with separate gate/up/down matrices."""

    hidden = _positive_int("hidden_size", hidden_size)
    intermediate = _positive_int("intermediate_size", intermediate_size)
    weights = 3 * hidden * intermediate
    biases = (2 * intermediate + hidden) if include_bias else 0
    return int(weights + biases)


def count_moe_parameters(spec: MoEModelSpec) -> ParameterBreakdown:
    """Count a model exactly under ``spec`` and expose every convention."""

    hidden = spec.hidden_size
    moe_layers = spec.layers - spec.dense_layers
    vocab_unit = spec.vocab_size * hidden
    vocab_total = spec.vocab_matrices_total * vocab_unit
    vocab_active = spec.active_vocab_matrices * vocab_unit
    attention = spec.attention_params_per_layer * spec.layers

    expert = expert_parameters(
        hidden_size=hidden,
        intermediate_size=spec.expert_intermediate,
        include_bias=spec.expert_bias,
    )
    expert_total = expert * (spec.routed_experts + spec.shared_experts) * moe_layers
    expert_active = expert * (spec.experts_per_token + spec.shared_experts) * moe_layers

    dense_unit = 3 * hidden * spec.dense_intermediate
    if spec.dense_bias:
        dense_unit += 2 * spec.dense_intermediate + hidden
    dense_ffn = dense_unit * spec.dense_layers

    router_per_layer = hidden * spec.routed_experts
    if spec.router_bias:
        router_per_layer += spec.routed_experts
    router = router_per_layer * moe_layers
    norms = 2 * hidden * spec.layers + hidden
    sinks = spec.sink_params_per_layer * spec.layers

    total = (
        vocab_total + attention + expert_total + dense_ffn + router + norms + sinks
        + spec.extra_total_parameters
    )
    active = (
        vocab_active + attention + expert_active + dense_ffn + router + norms + sinks
        + spec.extra_active_parameters
    )
    return ParameterBreakdown(
        total_parameters=int(total),
        active_parameters=int(active),
        vocab_total=int(vocab_total),
        vocab_active=int(vocab_active),
        attention=int(attention),
        routed_and_shared_experts_total=int(expert_total),
        routed_and_shared_experts_active=int(expert_active),
        dense_ffn=int(dense_ffn),
        router=int(router),
        norms=int(norms),
        sinks=int(sinks),
        extra_total=int(spec.extra_total_parameters),
        extra_active=int(spec.extra_active_parameters),
    )


def bytes_to_gib(byte_count: int) -> float:
    """Convert a non-negative byte count to binary gibibytes."""

    if isinstance(byte_count, bool) or not isinstance(byte_count, Integral) or byte_count < 0:
        raise ValueError("byte_count must be a non-negative integer")
    return int(byte_count) / 2**30


def expected_active_expert_fraction(*, experts: int, top_k: int, token_count: int) -> float:
    """Expected touched fraction under independent uniform token-choice routing."""

    expert_count = _positive_int("experts", experts)
    k = _positive_int("top_k", top_k)
    tokens = _positive_int("token_count", token_count)
    if k > expert_count:
        raise ValueError("top_k cannot exceed experts")
    return float(1.0 - (1.0 - k / expert_count) ** tokens)


def simulate_active_expert_fraction(
    *, experts: int, top_k: int, token_count: int, trials: int, seed: int
) -> float:
    """Monte Carlo check of the uniform independent-routing law."""

    expert_count = _positive_int("experts", experts)
    k = _positive_int("top_k", top_k)
    tokens = _positive_int("token_count", token_count)
    repeats = _positive_int("trials", trials)
    if k > expert_count:
        raise ValueError("top_k cannot exceed experts")
    if isinstance(seed, bool) or not isinstance(seed, Integral):
        raise TypeError("seed must be an integer")
    rng = np.random.default_rng(int(seed))
    fractions = []
    for _ in range(repeats):
        selected = np.stack([rng.choice(expert_count, size=k, replace=False) for _ in range(tokens)])
        fractions.append(np.unique(selected).size / expert_count)
    return float(np.mean(fractions))


def expert_parallel_payload_bytes(
    *, token_count: int, hidden_size: int, top_k: int, bytes_per_value: int,
    include_return: bool = True,
) -> int:
    """Idealized activation payload; metadata, padding, and topology are excluded."""

    tokens = _positive_int("token_count", token_count)
    hidden = _positive_int("hidden_size", hidden_size)
    k = _positive_int("top_k", top_k)
    width = _positive_int("bytes_per_value", bytes_per_value)
    one_way = tokens * hidden * k * width
    return one_way * (2 if include_return else 1)


models = REFERENCE["models"]
gqa = gqa_attention_parameters(hidden_size=2880, query_heads=64, kv_heads=8, head_width=64, include_bias=True)
assert gqa == models["gqa_params_per_layer_with_bias"]
mla = deepseek_v3_mla_attention_parameters()
assert mla == models["mla_params_per_layer"]
specs = {
    "oss120": MoEModelSpec(201_088, 2880, 36, gqa, 128, 4, 2880,
        vocab_matrices_total=2, active_vocab_matrices=1, router_bias=True,
        expert_bias=True, sink_params_per_layer=64),
    "oss20": MoEModelSpec(201_088, 2880, 24, gqa, 32, 4, 2880,
        vocab_matrices_total=2, active_vocab_matrices=1, router_bias=True,
        expert_bias=True, sink_params_per_layer=64),
    "deepseek_v3_reconstruction": MoEModelSpec(129_280, 7168, 61, mla, 256, 8, 2048,
        shared_experts=1, dense_layers=3, dense_intermediate=18_432,
        vocab_matrices_total=2, active_vocab_matrices=1, router_bias=True),
}
for name, spec in specs.items():
    counted = count_moe_parameters(spec)
    assert asdict(counted) == models[name]
    print(name, f"total={counted.total_parameters:,} active={counted.active_parameters:,} ratio={counted.total_to_active_ratio:.3f}x")
qwen = REFERENCE["qwen35_config"]
for source, target in [("hidden_size", "hidden_size"), ("num_hidden_layers", "layers"),
    ("num_experts", "routed_experts"), ("num_experts_per_tok", "selected_experts"),
    ("moe_intermediate_size", "routed_intermediate"),
    ("shared_expert_intermediate_size", "shared_intermediate"),
    ("mtp_num_hidden_layers", "mtp_layers")]:
    assert QWEN_TEXT[source] == qwen[target]
assert qwen == {"hidden_size": 2048, "layers": 40, "routed_experts": 256,
    "selected_experts": 8, "shared_experts": 1, "routed_intermediate": 512,
    "shared_intermediate": 512, "mtp_layers": 1}
assert math.comb(32, 8) == 10_518_300
assert math.comb(256, 8) == REFERENCE["combinations"]["C256_8_unconstrained_upper_bound"]
print('accounting and configs verified')

for key, row in REFERENCE["active_expert_law"].items():
    e, k, n = [int(part[1:]) for part in key.split("_")]
    theory = expected_active_expert_fraction(experts=e, top_k=k, token_count=n)
    assert math.isclose(theory, row["theory"], rel_tol=1e-12)
    assert abs(theory - row["simulation"]) < 0.01
    print(key, f"theory={theory:.4f} simulation={row['simulation']:.4f}")
latent = REFERENCE["latent_moe"]
assert latent["kimi_k3_hidden"] / latent["kimi_k3_latent"] == latent["payload_reduction"] == 2.0
comm = REFERENCE["expert_parallel"]
payload = expert_parallel_payload_bytes(token_count=comm["tokens"], hidden_size=comm["hidden_size"],
    top_k=comm["top_k"], bytes_per_value=comm["bytes_per_value"], include_return=True)
assert payload == comm["round_trip_bytes"] == 448 * 2**20
print("Kimi K3 latent reduction", latent["payload_reduction"], "x | EP payload", payload / 2**20, "MiB")
print('batch/LatentMoE/payload verified')

def expect_failure(exc, fn, *args, **kwargs):
    try:
        fn(*args, **kwargs)
    except exc:
        return
    raise AssertionError(f"expected {exc.__name__} from {fn.__name__}")

expect_failure(ValueError, route_logits, [[1., 2.]], top_k=3)
expect_failure(ValueError, route_logits, [[1., np.nan]], top_k=1)
expect_failure(ValueError, expert_capacity, token_count=8, top_k=2, expert_count=4, capacity_factor=0.)
expect_failure(ValueError, load_statistics, [0., 0.])
expect_failure(ValueError, aux_free_bias_step, [0., 0.], [1., 1.], rate=.1, mode="unknown")
expect_failure(ValueError, bytes_to_gib, -1)
expect_failure(ValueError, expert_parallel_payload_bytes, token_count=4, hidden_size=8, top_k=2, bytes_per_value=0)
expect_failure(ValueError, MoEModelSpec, 100, 64, 8, 1024, 4, 2, 128, dense_layers=2, dense_intermediate=0)
print('negative tests passed')

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_topk_router_explicit(
    logits: np.ndarray,
    top_k: int,
    *,
    selection_bias: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Minimal sigmoid router: bias changes selection, not clean mixture weights."""
    logits = np.asarray(logits, dtype=np.float64)
    if logits.ndim != 2 or not logits.size or not np.isfinite(logits).all():
        raise ValueError("expected non-empty finite logits[T,E]")
    if isinstance(top_k, (bool, np.bool_)) or not isinstance(top_k, (int, np.integer)) or not 1 <= top_k <= logits.shape[1]:
        raise ValueError("top_k must be an integer in [1,E]")
    log_scores = -np.logaddexp(0.0, -logits)
    clean = np.exp(log_scores)
    # Without bias, monotonic sigmoid permits ranking logits before saturation.
    selection = logits if selection_bias is None else clean + selection_bias
    if selection.shape != logits.shape or not np.isfinite(selection).all():
        raise ValueError("selection_bias must broadcast to [T,E] and be finite")
    indices = np.argsort(-selection, axis=-1, kind="stable")[:, :top_k]
    selected_log = np.take_along_axis(log_scores, indices, axis=-1)
    weights = np.exp(selected_log - selected_log.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)
    return indices, weights

def torch_topk_router_explicit(
    logits: torch.Tensor,
    top_k: int,
    *,
    selection_bias: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Minimal tensor counterpart using sigmoid, topk, and gather."""
    if logits.ndim != 2 or not logits.numel() or not logits.is_floating_point() or not torch.isfinite(logits).all():
        raise ValueError("expected non-empty finite floating logits[T,E]")
    if isinstance(top_k, bool) or not isinstance(top_k, int) or not 1 <= top_k <= logits.shape[1]:
        raise ValueError("top_k must be an integer in [1,E]")
    z = logits.float() if logits.dtype in (torch.float16, torch.bfloat16) else logits
    log_scores = -torch.logaddexp(torch.zeros_like(z), -z)
    selection = z if selection_bias is None else log_scores.exp() + selection_bias
    if selection.shape != logits.shape or not torch.isfinite(selection).all():
        raise ValueError("selection_bias must broadcast to [T,E] and be finite")
    _, indices = torch.topk(selection, top_k, dim=-1)
    selected_log = log_scores.gather(-1, indices)
    weights = selected_log.softmax(dim=-1)
    return indices, weights

logits_np = np.array([[2.1, -0.4, 0.8, 1.3], [0.2, 1.9, -1.1, 0.7]])
bias_np = np.array([0.0, 0.08, -0.03, 0.02])
i_np, w_np = numpy_topk_router_explicit(logits_np, 2, selection_bias=bias_np)
i_t, w_t = torch_topk_router_explicit(torch.tensor(logits_np), 2, selection_bias=torch.tensor(bias_np))
np.testing.assert_array_equal(i_np, i_t.numpy())
np.testing.assert_allclose(w_np, w_t.numpy())
extreme = np.array([[-1002., -1000., -1001., -1003.]])
with np.errstate(over="raise", invalid="raise", divide="raise"):
    ei, ew = numpy_topk_router_explicit(extreme, 2)
ti, tw = torch_topk_router_explicit(torch.tensor(extreme), 2)
np.testing.assert_array_equal(ei, [[1, 2]])
np.testing.assert_array_equal(ei, ti.numpy())
np.testing.assert_allclose(ew, tw.numpy(), rtol=1e-12, atol=1e-12)
z = torch.tensor(logits_np, requires_grad=True)
assert torch.autograd.gradcheck(lambda x: torch_topk_router_explicit(x, 2)[1], (z,))
one = torch_topk_router_explicit(z, 1)[1]
np.testing.assert_allclose(torch.autograd.grad(one.sum(), z)[0].numpy(), 0., atol=1e-12)
print("B09 routing, extreme logits, and selected-gate gradients: PASS")

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_expert_dispatch_explicit(
    hidden: np.ndarray,
    expert_indices: np.ndarray,
    gates: np.ndarray,
    w_in: np.ndarray,
    w_out: np.ndarray,
) -> np.ndarray:
    """Minimal token/slot loop that makes dispatch and weighted combine visible."""
    if expert_indices.ndim != 2 or expert_indices.shape != gates.shape or expert_indices.shape[0] != hidden.shape[0]:
        raise ValueError("indices and gates must have matching shape [T,K]")
    if not np.issubdtype(expert_indices.dtype, np.integer):
        raise ValueError("expert indices must be integers")
    if np.any(expert_indices < 0) or np.any(expert_indices >= w_in.shape[0]):
        raise ValueError("expert index out of range")
    output = np.zeros_like(hidden)
    for token in range(hidden.shape[0]):
        for slot in range(expert_indices.shape[1]):
            expert = int(expert_indices[token, slot])
            mid = np.tanh(hidden[token] @ w_in[expert])
            output[token] += gates[token, slot] * (mid @ w_out[expert])
    return output

def torch_expert_dispatch_explicit(
    hidden: torch.Tensor,
    expert_indices: torch.Tensor,
    gates: torch.Tensor,
    w_in: torch.Tensor,
    w_out: torch.Tensor,
) -> torch.Tensor:
    """Group assignments by expert and combine them with index_add_."""
    if expert_indices.ndim != 2 or expert_indices.shape != gates.shape or expert_indices.shape[0] != hidden.shape[0]:
        raise ValueError("indices and gates must have matching shape [T,K]")
    if expert_indices.dtype not in (torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8):
        raise ValueError("expert indices must be integers")
    if torch.any(expert_indices < 0) or torch.any(expert_indices >= w_in.shape[0]):
        raise ValueError("expert index out of range")
    output = torch.zeros_like(hidden)
    for expert in range(w_in.shape[0]):
        token_ids, slots = torch.where(expert_indices == expert)
        if token_ids.numel() == 0:
            continue
        mid = torch.tanh(hidden[token_ids] @ w_in[expert])
        expert_out = (mid @ w_out[expert]) * gates[token_ids, slots, None]
        output.index_add_(0, token_ids, expert_out)
    return output

rng = np.random.default_rng(10)
x_np = rng.standard_normal((6, 4))
idx_np = np.array([[0,1],[1,2],[2,0],[0,2],[1,0],[2,1]])
gates_np = np.full((6,2), 0.5)
w1_np = rng.standard_normal((3,4,5)); w2_np = rng.standard_normal((3,5,4))
np_out = numpy_expert_dispatch_explicit(x_np, idx_np, gates_np, w1_np, w2_np)
t_out = torch_expert_dispatch_explicit(torch.tensor(x_np), torch.tensor(idx_np), torch.tensor(gates_np), torch.tensor(w1_np), torch.tensor(w2_np))
np.testing.assert_allclose(np_out, t_out.numpy(), rtol=1e-10, atol=1e-10)
args = [torch.tensor(x_np, requires_grad=True), torch.tensor(idx_np),
        torch.tensor(gates_np, requires_grad=True), torch.tensor(w1_np, requires_grad=True),
        torch.tensor(w2_np, requires_grad=True)]
assert torch.autograd.gradcheck(
    lambda x, g, a, b: torch_expert_dispatch_explicit(x, args[1], g, a, b),
    (args[0], args[2], args[3], args[4]))
for invalid in (-1, w1_np.shape[0], .5, np.nan, True):
    bad = np.full(idx_np.shape, invalid)
    for dispatch, values in (
        (numpy_expert_dispatch_explicit, (x_np,bad,gates_np,w1_np,w2_np)),
        (torch_expert_dispatch_explicit, tuple(torch.tensor(a) for a in (x_np,bad,gates_np,w1_np,w2_np))),
    ):
        try: dispatch(*values)
        except ValueError: pass
        else: raise AssertionError("invalid expert index accepted")
duplicate = np.zeros_like(idx_np)
expected = idx_np.shape[1] * (np.tanh(x_np @ w1_np[0]) @ w2_np[0])
actual = numpy_expert_dispatch_explicit(x_np,duplicate,np.ones_like(gates_np),w1_np,w2_np)
np.testing.assert_allclose(actual,expected)
actual_t = torch_expert_dispatch_explicit(torch.tensor(x_np),torch.tensor(duplicate),
    torch.ones_like(torch.tensor(gates_np)),torch.tensor(w1_np),torch.tensor(w2_np))
np.testing.assert_allclose(actual_t.numpy(),expected)
print("B10 values, gradients, repeated assignments, and invalid indices: PASS")

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})

ru = DEMO_LANGUAGE == "ru"
rng = np.random.default_rng(SEED)
cluster = np.repeat(np.arange(3), 300)
features = rng.normal(size=(900, 8)) + np.eye(3, 8)[cluster] * 2
router_w = rng.normal(size=(8, 6))
logits = features @ router_w
log_scores = -np.logaddexp(0., -logits)
scores = np.exp(log_scores)

def route(bias):
    indices = np.argsort(-(scores + bias), axis=1, kind="stable")[:, :2]
    selected = np.take_along_axis(log_scores, indices, axis=1)
    gates = np.exp(selected - selected.max(axis=1, keepdims=True))
    gates /= gates.sum(axis=1, keepdims=True)
    return indices, gates, np.bincount(indices.ravel(), minlength=6)

bias = np.zeros(6)
history = []
for step in range(101):
    top2, gates, counts = route(bias)
    history.append(counts.std() / counts.mean())
    if step < 100:
        bias += .01 * np.sign(counts.mean() - counts)
assert history[-1] < history[0]
assert counts.sum() == 1800
print("CV:", round(history[0], 4), "→", round(history[-1], 4))
fig, axes = plt.subplots(1, 2, figsize=(11, 4.8), layout="constrained")
axes[0].plot(history, color="#5E81AC")
axes[0].set(xlabel="Обновления смещений" if ru else "Bias updates", ylabel="CV",
            title="Фиксированные оценки" if ru else "Fixed router scores")
heat = np.zeros((6, 3), dtype=int)
for c, selected in zip(cluster, top2):
    heat[selected, c] += 1
im = axes[1].imshow(heat / 600, cmap="Blues", vmin=0, vmax=1, aspect="auto")
axes[1].set(xlabel="Синтетический кластер" if ru else "Synthetic cluster",
            ylabel="Эксперт" if ru else "Expert", xticks=[0,1,2],
            title="Доля назначений в кластере" if ru else "Assignment share within cluster")
fig.colorbar(im, ax=axes[1])
assert np.allclose((heat / 600).sum(axis=0), 1.)
plt.show()

# Count co-activation separately from the effect of changing expert outputs.
co = np.zeros((6, 6), dtype=int)
for a, b in top2:
    co[a, b] += 1
    co[b, a] += 1
assert np.all(co == co.T) and np.all(np.diag(co) == 0)
assert np.array_equal(co.sum(axis=1), counts)  # exactly one partner at top-2
conditional = co / counts[:, None]
assert np.allclose(conditional.sum(axis=1), 1.)
w1 = rng.normal(scale=.3, size=(6, 8, 12))
w2 = rng.normal(scale=.3, size=(6, 12, 8))
expert_outputs = np.stack([np.tanh(features @ a) @ b for a,b in zip(w1,w2)], axis=1)
selected_outputs = expert_outputs[np.arange(len(features))[:, None], top2]
original = (gates[..., None] * selected_outputs).sum(axis=1)
effects = []
for expert in range(6):
    # Zero an expert contribution while keeping routes and all other gates fixed.
    kept = gates * (top2 != expert)
    intervened = (kept[..., None] * selected_outputs).sum(axis=1)
    effects.append(np.mean((intervened - original)**2))
assert np.all(np.isfinite(effects)) and min(effects) >= 0
fig, axes = plt.subplots(1,2,figsize=(11,4.8),layout="constrained")
im=axes[0].imshow(conditional,cmap="Blues",vmin=0,vmax=1)
axes[0].set(xlabel="Партнёр" if ru else "Partner expert",ylabel="Выбранный эксперт" if ru else "Selected expert",
    title="P(партнёр | эксперт)" if ru else "P(partner | expert)")
fig.colorbar(im,ax=axes[0])
axes[1].bar(range(6),effects,color="#D08770")
axes[1].set(xlabel="Обнулённый эксперт" if ru else "Zeroed expert",
    ylabel="MSE выхода" if ru else "Output MSE",
    title="При неизменных маршрутах" if ru else "With routes held fixed")
plt.show()
print("Это изменение вектора, не падение качества задачи." if ru else "This measures vector change, not task-quality loss.")