/**
 * moe_model.ts — SharedExpertMoE: a shared-expert hybrid Mixture-of-Experts FFN.
 *
 * The sparsity design behind Evermind's generator. Each token is processed by:
 *   • a DENSE shared expert that is ALWAYS active (carries continuous learning;
 *     the part the online-distillation signal flows into), plus
 *   • the top-k of N routed experts, gated by a learned router and combined by a
 *     softmax over the selected experts.
 *
 *   y = SharedFFN(x) + Σ_{e ∈ topk(x)} gate_e · Expert_e(x)
 *
 * This is the DeepSeekMoE "shared-expert isolation" pattern: the dense backbone
 * resolves the online-learning attribution problem (you distil into ONE always-on
 * path), while the routed experts add web-pageable capacity (each expert's
 * weights are an independent checkpoint — see {@link SharedExpertMoE.exportExpert}
 * — so a host can stream only the experts a token activates).
 *
 * Pure-TS CPU reference (Float32Array, exact forward + backward), mirroring
 * {@link LimbicModel}'s WebGPU-or-fallback contract — the WGSL kernel path
 * (router gate + expert FFN GEMM) is a numerically-identical future acceleration.
 *
 * Activation is ReLU for an exact, unambiguous gradient in the reference path;
 * production may swap GELU/SwiGLU behind the same shapes.
 */
export interface MoEConfig {
    /** Model (token) dimension — FFN input/output width. Default 64. */
    modelDim: number;
    /** Hidden width of each expert FFN. Default 128. */
    hiddenDim: number;
    /** Number of routed experts. Default 8. */
    numExperts: number;
    /** Experts activated per token (top-k). Default 2. Must be ≤ numExperts. */
    topK: number;
    /** Deterministic init seed for reproducible cold-start weights. */
    seed?: number;
}
export declare const DEFAULT_MOE_CONFIG: Required<Omit<MoEConfig, "seed">>;
/** Fixed default init seed — reproducible byte-identical cold start across machines. */
export declare const DEFAULT_MOE_SEED = 1299137793;
/** A named trainable parameter tensor (flat row-major). */
export interface MoEParam {
    name: string;
    data: Float32Array;
    numel: number;
}
/** Result of routing a token: which experts fire and with what combine weights. */
export interface RouteResult {
    /** Indices of the selected top-k experts, highest router logit first. */
    experts: number[];
    /** Combine weights (softmax over the selected logits), index-aligned to `experts`. */
    gates: number[];
    /** Full softmax over ALL experts — the load-balancing signal. */
    probs: Float32Array;
}
/** Per-token forward intermediates retained for the backward pass. */
interface MoECache {
    x: Float32Array;
    route: RouteResult;
    sharedPre: Float32Array;
    sharedH: Float32Array;
    expertOut: Float32Array[];
    expertPre: Float32Array[];
    expertH: Float32Array[];
}
/**
 * Accumulates router statistics over a batch to compute the load-balancing
 * auxiliary loss `E · Σ_e f_e · P_e` (Switch/GShard). Minimised (→ near 1) when
 * dispatch is uniform; large (→ near E) when the router collapses onto few
 * experts. Add it to the task loss with a small coefficient to keep experts busy.
 */
export declare class LoadBalanceAccumulator {
    private readonly numExperts;
    private readonly counts;
    private readonly probSum;
    private tokens;
    constructor(numExperts: number);
    observe(route: RouteResult): void;
    /** The load-balance loss over everything observed so far (0 if no tokens). */
    loss(): number;
}
export declare class SharedExpertMoE {
    readonly config: Required<Omit<MoEConfig, "seed">>;
    /** Router weights: numExperts × modelDim (no bias). */
    wr: Float32Array;
    private gWr;
    private readonly shared;
    private readonly experts;
    constructor(config?: Partial<MoEConfig>);
    /** Route a token: router logits → top-k → combine gates + full softmax probs. */
    route(x: Float32Array): RouteResult;
    /** Forward a single token. Returns the output and a cache for {@link backward}. */
    forward(input: ArrayLike<number>): {
        output: Float32Array;
        route: RouteResult;
        cache: MoECache;
    };
    /**
     * Accumulate gradients for one token given dL/d(output). Trains the shared
     * expert, the selected routed experts, and the router (so it learns to weight
     * the experts that reduce loss). Call {@link zeroGrad} before a batch and apply
     * an optimiser after. Load balancing is a separate signal (see
     * {@link LoadBalanceAccumulator}). Returns dL/d(input) so the FFN can stack
     * inside a residual block (e.g. {@link EvermindLM}).
     */
    backward(dOutput: ArrayLike<number>, cache: MoECache): Float32Array;
    /**
     * Add the load-balancing auxiliary-loss gradient for one token into the router
     * gradient. `L_aux = E·Σ_e f_e·P̄_e` (Switch/GShard); `f` (per-batch dispatch
     * fractions) is treated as a stop-grad constant, so only the full softmax `P`
     * carries gradient: ∂L_aux/∂logit_j = scale·P_j·(f_j − Σ_e f_e·P_e), where the
     * caller passes `scale = auxWeight·E/T`. Keeps the router from collapsing onto a
     * few experts. Call once per token over the batch, after {@link backward}.
     */
    auxGradStep(x: Float32Array, probs: Float32Array, f: Float32Array, scale: number): void;
    /** All trainable parameters in canonical order: router, shared, then experts. */
    parameters(): MoEParam[];
    /** Gradient buffers, index-aligned with {@link parameters}. */
    gradients(): MoEParam[];
    zeroGrad(): void;
    /** One routed expert's weights as a standalone checkpoint (the web-paging unit). */
    exportExpert(index: number): MoEParam[];
    /**
     * Serialise all weights to a compact "MoE0" binary. Layout: magic, version,
     * [modelDim, hiddenDim, numExperts, topK], then params in {@link parameters}
     * order. fp16 (v2) halves the size; f32 (v1) is exact.
     */
    exportWeights(opts?: {
        fp16?: boolean;
    }): ArrayBuffer;
    /** Load weights from an "MoE0" binary. Validates magic + dims. */
    loadWeights(buffer: ArrayBuffer): void;
}
export {};
//# sourceMappingURL=moe_model.d.ts.map