@seanhogg/builderforce-memory-engine 2026.9.28 → 2026.9.30

package/src/experience/types.ts

@@ -0,0 +1,162 @@
+/**
+ * experience/types.ts — what Evermind has DONE, as opposed to what it knows.
+ *
+ * Facts live in `@seanhogg/builderforce-memory` (MemoryStore / EvermindCognition) and
+ * weights in the model; this is the third kind of learned item: experience.
+ *   - an **Episode** is one demonstration — a person doing a task in a program, recorded
+ *     as semantic steps (the UI element acted on and the value set, never raw keystrokes);
+ *   - a **Skill** is what Train Once compiles from an episode — a procedure that runs
+ *     again with new values, with approval gates on irreversible steps;
+ *   - a **Run** is one execution of a skill, with its step-by-step audit trail.
+ *
+ * Plain data, JSON-shaped (camelCase, actions tagged by `kind`), so any recorder — the
+ * native Windows UI Automation one in Synapse included — produces it and any runtime
+ * persists it. Timestamps are epoch milliseconds.
+ */
+
+/** A UI element, identified the way the platform's accessibility API sees it. */
+export interface ElementRef {
+  name: string;
+  automationId: string;
+  /** `Button`, `Edit`, `MenuItem`, … */
+  controlType: string;
+  className: string;
+  /** The top-level window the element sits in. */
+  windowTitle: string;
+  processName: string;
+  /** Where it was acted on, relative to the window's top-left — the replay fallback. */
+  relX: number;
+  relY: number;
+}
+
+/** One thing the person did. A secret field records that it was filled, never the value. */
+export type Action =
+  | { kind: 'click'; target: ElementRef }
+  | { kind: 'setValue'; target: ElementRef; value: string; secret?: boolean }
+  | { kind: 'keys'; keys: string; target?: ElementRef | null };
+
+export interface RecordedStep {
+  id: string;
+  /** Milliseconds since the recording started. */
+  atMs: number;
+  action: Action;
+  /** Screenshot file name inside the episode's folder, if one was taken. */
+  screenshot?: string | null;
+}
+
+export interface Episode {
+  id: string;
+  name: string;
+  program: string;
+  args: string[];
+  startedAt: number;
+  endedAt?: number | null;
+  steps: RecordedStep[];
+}
+
+export interface EpisodeSummary {
+  id: string;
+  name: string;
+  program: string;
+  startedAt: number;
+  steps: number;
+}
+
+/** A value a skill asks for at run time. */
+export interface SkillParam {
+  /** Stable key, e.g. `amount`. */
+  name: string;
+  /** What the field was called where it was typed, e.g. `Amount`. */
+  label: string;
+  /** The value from the demonstration; never set for a secret. */
+  default?: string | null;
+  secret: boolean;
+}
+
+export type SkillValue = { from: 'param'; name: string } | { from: 'literal'; text: string };
+
+export type SkillAction =
+  | { kind: 'click'; target: ElementRef }
+  | { kind: 'setValue'; target: ElementRef; value: SkillValue }
+  | { kind: 'keys'; keys: string; target?: ElementRef | null };
+
+export interface SkillStep {
+  id: string;
+  action: SkillAction;
+  /** Pause and ask the person before this step. */
+  requiresApproval: boolean;
+  screenshot?: string | null;
+}
+
+/** When a skill runs by itself. Local time of day for `day`. */
+export type Schedule = { every: 'minutes'; minutes: number } | { every: 'day'; hour: number; minute: number };
+
+/** A skill's routine: at most one per skill, so it lives on the skill. */
+export interface SkillRoutine {
+  schedule: Schedule;
+  /** Values for the non-secret parameters; secrets come from the host's vault. */
+  values: Record<string, string>;
+  enabled: boolean;
+  lastRunAt?: number | null;
+}
+
+export interface Skill {
+  id: string;
+  name: string;
+  program: string;
+  args: string[];
+  params: SkillParam[];
+  steps: SkillStep[];
+  sourceEpisode: string;
+  createdAt: number;
+  routine?: SkillRoutine | null;
+}
+
+export type RunStatus = 'running' | 'succeeded' | 'failed' | 'stopped' | 'denied';
+
+export type StepOutcome = 'ok' | 'fallback' | 'failed' | 'approved' | 'denied';
+
+/** One line of a run's audit trail. A step can have two (approved, then ok). */
+export interface RunStepLog {
+  idx: number;
+  stepId: string;
+  outcome: StepOutcome;
+  detail?: string | null;
+  at: number;
+}
+
+export interface Run {
+  id: string;
+  skillId: string;
+  skillName: string;
+  /** `manual` or `routine`. */
+  trigger: string;
+  status: RunStatus;
+  startedAt: number;
+  endedAt?: number | null;
+  error?: string | null;
+  steps: RunStepLog[];
+}
+
+/**
+ * One adaptation of the private model on experience: what it learned and how the fit
+ * went. Measured when the adaptation ran, never estimated afterwards.
+ */
+export interface Adaptation {
+  /** N of the `+expN` version this adaptation wrote. */
+  index: number;
+  version: string;
+  /** Episode and skill ids it learned. */
+  learned: string[];
+  passes: number;
+  /** Mean fit loss over the passes. */
+  loss: number;
+  at: number;
+}
+
+/** A short random id — unique enough for one person's store. */
+export function newExperienceId(): string {
+  const bytes = new Uint8Array(8);
+  globalThis.crypto.getRandomValues(bytes);
+  return Array.from(bytes, (b) => b.toString(16).padStart(2, '0')).join('');
+}

package/src/export/onnx.ts

@@ -95,13 +95,6 @@ class GraphBuilder {
   }
 }
 
-/** Transpose a 2-D row-major matrix [rows, cols] → [cols, rows]. */
-function transpose2d(data: Float32Array, rows: number, cols: number): Float32Array {
-  const out = new Float32Array(data.length);
-  for (let r = 0; r < rows; r++) for (let c = 0; c < cols; c++) out[c * rows + r] = data[r * cols + c]!;
-  return out;
-}
-
 /** RMSNorm(x, gain) as graph ops. Returns the normalised value name. */
 function rmsNorm(g: GraphBuilder, x: string, gainInit: string, eps: string, axes: string): string {
   const sq = g.op("Mul", [x, x], "rms_sq");

package/src/import/foreign/falcon_mamba.ts

@@ -1,6 +1,9 @@
 /**
- * import/foreign/falcon_mamba.ts — weight port for Falcon-Mamba
- * (`FalconMambaForCausalLM`, `model_type: "falcon_mamba"`, Mamba-1 / S6).
+ * import/foreign/falcon_mamba.ts — weight port for the Mamba-1 / S6 family:
+ * Falcon-Mamba (`FalconMambaForCausalLM`, `model_type: "falcon_mamba"`) and the
+ * original Mamba (`MambaForCausalLM`, `model_type: "mamba"` — `state-spaces/mamba-*`).
+ * They share one tensor layout and one config vocabulary, so one rule table and one
+ * `describe` serve both ({@link mamba1Adapter}).
  *
  * Source layout confirmed against `tiiuae/falcon-mamba-7b`'s published
  * `config.json`, `model.safetensors.index.json` (643 tensors = 10 per layer × 64
@@ -34,6 +37,8 @@
  *   • Linear weights are `[out, in]` row-major on both sides (no transpose).
  *
  * Fidelity caveats recorded on the plan rather than hidden:
+ *   • (Falcon-Mamba only) the dt/B/C norms below. Plain Mamba has none, so its port
+ *     is exact.
  *   • `use_bias: false` upstream, so `in_proj` / `x_proj` / `out_proj` ship no
  *     bias; this engine's blocks always own one, and they are filled with zeros
  *     (an exact no-op, not an approximation).
@@ -59,8 +64,6 @@ import {
   type PortTarget,
 } from "./adapter.js";
 
-const ID = "falcon_mamba";
-
 /** `backbone.embeddings` (HF) vs `backbone.embedding` (Mamba-SSM native). */
 const EMBEDDING_NAMES = ["backbone.embeddings.weight", "backbone.embedding.weight"];
 
@@ -131,43 +134,63 @@ function rules(a: {
   return out;
 }
 
-export const falconMambaAdapter: ForeignMambaAdapter = {
-  id: ID,
-  label: "Falcon-Mamba (Mamba-1 / S6)",
-  modelTypes: ["falcon_mamba"],
-  architectures: ["FalconMambaForCausalLM"],
+/** A Mamba-1 family adapter: the shared rule table under one checkpoint family's identity. */
+export function mamba1Adapter(identity: {
+  id: string;
+  label: string;
+  modelTypes: readonly string[];
+  architectures: readonly string[];
+}): ForeignMambaAdapter {
+  const ID = identity.id;
+  return {
+    ...identity,
+
+    describe(config: ForeignConfig): PortTarget {
+      const dModel = requireInt(config, "hidden_size");
+      const numLayers = requireInt(config, "num_hidden_layers");
+      const vocabSize = requireInt(config, "vocab_size");
+      const dState = requireInt(config, "state_size");
+      const dConv = requireInt(config, "conv_kernel");
+      const dInner = readInnerSize(config, dModel);
+      const expand = expandFactor(dInner, dModel, identity.label);
+      const dtRank = readTimeStepRank(config, dModel);
 
-  describe(config: ForeignConfig): PortTarget {
-    const dModel = requireInt(config, "hidden_size");
-    const numLayers = requireInt(config, "num_hidden_layers");
-    const vocabSize = requireInt(config, "vocab_size");
-    const dState = requireInt(config, "state_size");
-    const dConv = requireInt(config, "conv_kernel");
-    const dInner = readInnerSize(config, dModel);
-    const expand = expandFactor(dInner, dModel, "Falcon-Mamba");
-    const dtRank = readTimeStepRank(config, dModel);
+      const modelConfig: HybridMambaModelConfig = {
+        vocabSize,
+        dModel,
+        numLayers,
+        dState,
+        dConv,
+        expand,
+        // Every layer is a Mamba-1 mixer (the model's default schedule), but the
+        // dt rank and conv bias are per-block settings the shorthand can't carry.
+        defaultMamba1: { dtRank, biasConv: true },
+      };
 
-    const modelConfig: HybridMambaModelConfig = {
-      vocabSize,
-      dModel,
-      numLayers,
-      dState,
-      dConv,
-      expand,
-      // Every layer is a Mamba-1 mixer (the model's default schedule), but the
-      // dt rank and conv bias are per-block settings the shorthand can't carry.
-      defaultMamba1: { dtRank, biasConv: true },
-    };
+      const plan: PortPlan = {
+        adapter: ID,
+        rules: rules({ vocabSize, dModel, numLayers, dInner, dState, dConv, dtRank }),
+        // Nothing is partially consumed: every Mamba-1 tensor this plan
+        // touches is taken whole. `lm_head.weight` is untouched entirely and
+        // surfaces through `unmappedSources`.
+        discards: [],
+      };
 
-    const plan: PortPlan = {
-      adapter: ID,
-      rules: rules({ vocabSize, dModel, numLayers, dInner, dState, dConv, dtRank }),
-      // Nothing is partially consumed: every Falcon-Mamba tensor this plan
-      // touches is taken whole. `lm_head.weight` is untouched entirely and
-      // surfaces through `unmappedSources`.
-      discards: [],
-    };
+      return { modelConfig, plan };
+    },
+  };
+}
+
+export const falconMambaAdapter = mamba1Adapter({
+  id: "falcon_mamba",
+  label: "Falcon-Mamba (Mamba-1 / S6)",
+  modelTypes: ["falcon_mamba"],
+  architectures: ["FalconMambaForCausalLM"],
+});
 
-    return { modelConfig, plan };
-  },
-};
+export const mambaAdapter = mamba1Adapter({
+  id: "mamba",
+  label: "Mamba (Mamba-1 / S6)",
+  modelTypes: ["mamba"],
+  architectures: ["MambaForCausalLM"],
+});

package/src/import/foreign/index.ts

@@ -6,7 +6,7 @@
  * to build plus a name-keyed set of weights, together with an explicit account of
  * anything that did not map cleanly.
  *
- * Supported: Falcon-Mamba (Mamba-1) and Codestral-Mamba (Mamba-2). Transformer
+ * Supported: Mamba and Falcon-Mamba (Mamba-1) and Codestral-Mamba (Mamba-2). Transformer
  * checkpoints are NOT portable at any tensor naming — distillation is the only
  * route, and {@link foreignMambaAdapterFor} says so when it rejects one.
  */
@@ -33,5 +33,8 @@ export type {
 
 export type { ForeignConfig, ForeignMambaAdapter, PortTarget } from "./adapter.js";
 
-export { falconMambaAdapter } from "./falcon_mamba.js";
+export { falconMambaAdapter, mambaAdapter, mamba1Adapter } from "./falcon_mamba.js";
+export { mambaConfigFromTensors } from "./infer.js";
+export { resizePortedVocab } from "./vocab.js";
+export type { ResizeVocabOptions } from "./vocab.js";
 export { codestralMambaAdapter } from "./codestral_mamba.js";

package/src/import/foreign/infer.ts

@@ -0,0 +1,62 @@
+/**
+ * import/foreign/infer.ts — a Mamba-1 checkpoint's config, read off its tensors.
+ *
+ * `state-spaces/mamba-*` repositories publish `model.safetensors` with a
+ * `config.json` in the original Mamba-SSM vocabulary (`d_model`, `n_layer`, …)
+ * that no adapter reads, and a dropped safetensors file often arrives with no
+ * config at all. Every geometry fact the port needs is in the tensor shapes, so
+ * this derives the Hugging Face `MambaConfig` the {@link mambaAdapter} expects
+ * instead of anyone restating the derivation.
+ */
+
+import type { NamedTensor } from "../../export/tensors.js";
+import type { ForeignConfig } from "./adapter.js";
+import { normaliseSourceName } from "./plan.js";
+
+function fail(detail: string): never {
+  throw new Error(`import/foreign: cannot read a Mamba-1 config from these tensors — ${detail}`);
+}
+
+/**
+ * The `MambaForCausalLM` config implied by a Mamba-1 checkpoint's tensors:
+ * embedding `[vocab, d_model]`, per-layer `in_proj [2·d_inner, d_model]`,
+ * `conv1d [d_inner, (1,) d_conv]`, `A_log [d_inner, d_state]` and
+ * `x_proj [dt_rank + 2·d_state, d_inner]`.
+ */
+export function mambaConfigFromTensors(tensors: readonly NamedTensor[]): ForeignConfig {
+  const byName = new Map(tensors.map((t) => [normaliseSourceName(t.name), t]));
+  const shape = (name: string) => byName.get(name)?.shape;
+
+  const emb = shape("backbone.embeddings.weight") ?? shape("backbone.embedding.weight");
+  if (!emb || emb.length !== 2) fail("no backbone.embedding(s).weight [vocab, d_model]");
+  const [vocabSize, dModel] = emb as [number, number];
+
+  let numLayers = 0;
+  while (byName.has(`backbone.layers.${numLayers}.mixer.in_proj.weight`)) numLayers++;
+  if (numLayers === 0) fail("no backbone.layers.N.mixer.in_proj.weight");
+
+  const m = "backbone.layers.0.mixer";
+  const inProj = shape(`${m}.in_proj.weight`)!;
+  const conv = shape(`${m}.conv1d.weight`);
+  const aLog = shape(`${m}.A_log`);
+  const xProj = shape(`${m}.x_proj.weight`);
+  if (!conv || !aLog || !xProj) fail(`layer 0 is missing conv1d.weight, A_log or x_proj.weight`);
+
+  const dInner = inProj[0]! / 2;
+  const dState = aLog[1]!;
+  const dConv = conv[conv.length - 1]!;
+  const dtRank = xProj[0]! - 2 * dState;
+  if (!Number.isInteger(dInner) || dtRank <= 0) fail(`inconsistent shapes (in_proj ${inProj.join("×")}, x_proj ${xProj.join("×")})`);
+
+  return {
+    model_type: "mamba",
+    architectures: ["MambaForCausalLM"],
+    vocab_size: vocabSize,
+    hidden_size: dModel,
+    num_hidden_layers: numLayers,
+    intermediate_size: dInner,
+    state_size: dState,
+    conv_kernel: dConv,
+    time_step_rank: dtRank,
+  };
+}

package/src/import/foreign/registry.ts

@@ -12,12 +12,13 @@ import type { NamedTensor } from "../../export/tensors.js";
 import type { HybridMambaModelConfig } from "../../model/mamba_model.js";
 import { safetensorsToTensors } from "../safetensors.js";
 import { executePortPlan, type PortDiscard, type PortedTensors, type SynthesisedTarget } from "./plan.js";
-import { falconMambaAdapter } from "./falcon_mamba.js";
+import { falconMambaAdapter, mambaAdapter } from "./falcon_mamba.js";
 import { codestralMambaAdapter } from "./codestral_mamba.js";
 import type { ForeignConfig, ForeignMambaAdapter } from "./adapter.js";
 
 /** Every foreign SSM architecture this engine can warm-start from. */
 export const FOREIGN_MAMBA_ADAPTERS: readonly ForeignMambaAdapter[] = [
+  mambaAdapter,
   falconMambaAdapter,
   codestralMambaAdapter,
 ];

package/src/import/foreign/vocab.ts

@@ -0,0 +1,50 @@
+/**
+ * import/foreign/vocab.ts — fit a ported checkpoint to a different vocabulary size.
+ *
+ * A ported model keeps its source tokenizer's vocabulary. To pair it with another
+ * tokenizer the embedding table (tied to the LM head) must have that tokenizer's row
+ * count: surplus rows are dropped, missing rows are drawn small and random so the
+ * new tokens start near the embedding cloud's centre. The new rows are recorded on
+ * the port as synthesised, never passed off as ported.
+ */
+
+import { SeededRng, gaussianArray } from "../../utils/rng.js";
+import type { PortedCheckpoint } from "./registry.js";
+
+export interface ResizeVocabOptions {
+  /** Standard deviation of new rows. Default 0.02, the usual embedding init. */
+  std?: number;
+  /** Seed for the new rows, so the same resize produces the same checkpoint. Default 1. */
+  seed?: number;
+}
+
+/** `ported` with an embedding of `vocabSize` rows and a model config to match. */
+export function resizePortedVocab(ported: PortedCheckpoint, vocabSize: number, opts: ResizeVocabOptions = {}): PortedCheckpoint {
+  if (!Number.isInteger(vocabSize) || vocabSize <= 0) {
+    throw new Error(`import/foreign: vocabulary size must be a positive integer, got ${vocabSize}`);
+  }
+  const from = ported.modelConfig.vocabSize;
+  if (vocabSize === from) return ported;
+
+  const emb = ported.weights.get("embedding");
+  if (!emb) throw new Error('import/foreign: the port has no "embedding" to resize');
+  const dModel = ported.modelConfig.dModel;
+
+  const resized = new Float32Array(vocabSize * dModel);
+  resized.set(emb.subarray(0, Math.min(from, vocabSize) * dModel));
+  const synthesisedTargets = [...ported.synthesisedTargets];
+  if (vocabSize > from) {
+    const std = opts.std ?? 0.02;
+    const rng = new SeededRng(opts.seed ?? 1);
+    resized.set(gaussianArray((vocabSize - from) * dModel, std, () => rng.next()), from * dModel);
+    synthesisedTargets.push({
+      target: "embedding",
+      value: std,
+      reason: `rows ${from}..${vocabSize - 1} drawn from N(0, ${std}²) to reach a ${vocabSize}-token vocabulary`,
+    });
+  }
+
+  const weights = new Map(ported.weights);
+  weights.set("embedding", resized);
+  return { ...ported, weights, synthesisedTargets, modelConfig: { ...ported.modelConfig, vocabSize } };
+}

package/src/import/index.ts

@@ -25,7 +25,11 @@ export {
   executePortPlan,
   normaliseSourceName,
   falconMambaAdapter,
+  mambaAdapter,
+  mamba1Adapter,
   codestralMambaAdapter,
+  mambaConfigFromTensors,
+  resizePortedVocab,
 } from "./foreign/index.js";
 export type {
   ForeignConfig,
@@ -38,4 +42,5 @@ export type {
   PortedCheckpoint,
   PortedTensors,
   SynthesisedTarget,
+  ResizeVocabOptions,
 } from "./foreign/index.js";

package/src/import/safetensors.ts

@@ -4,7 +4,7 @@
  * The exact inverse of {@link ../export/safetensors.ts}: parse the 8-byte
  * little-endian u64 header length, the JSON header (name → {dtype, shape,
  * data_offsets}), then decode each tensor's raw bytes into a Float32Array.
- * F32 and F16 are supported (F16 is widened to f32 on read). This is the reader
+ * F32, F16 and BF16 are supported (both halves are widened to f32 on read). This is the reader
  * half of the warm-start / weight-port path — export was one-directional before.
  */
 
@@ -63,5 +63,12 @@ function decodeTensor(seg: Uint8Array, dtype: string, numel: number, name: strin
     for (let i = 0; i < numel; i++) data[i] = fp16ToFloat(dv.getUint16(i * 2, true));
     return data;
   }
-  throw new Error(`safetensors: unsupported dtype "${dtype}" for "${name}" (expected F32 or F16)`);
+  if (d === "BF16" || d === "BFLOAT16") {
+    // bfloat16 is the top half of an f32: shift it back into place.
+    if (seg.length < numel * 2) throw new Error(`safetensors: "${name}" BF16 bytes ${seg.length} < ${numel * 2}`);
+    const bits = new Uint32Array(data.buffer);
+    for (let i = 0; i < numel; i++) bits[i] = dv.getUint16(i * 2, true) << 16;
+    return data;
+  }
+  throw new Error(`safetensors: unsupported dtype "${dtype}" for "${name}" (expected F32, F16 or BF16)`);
 }

package/src/index.ts

@@ -35,7 +35,9 @@ export type {
 // ── Evermind (the .evermind artifact + the shared learning pipeline) ───────────
 export {
     EvermindModelPackage,
+    PACKAGE_HEADER_BYTES,
     adaptAndDiff,
+    adaptPackage,
     EVERMIND_ADAPT_MAX_CHARS,
     parseDeltaLearnPayload,
     buildDeltaLearnPayload,
@@ -53,6 +55,7 @@ export type {
     PackageMeta,
     ValidationResult,
     AdaptResult,
+    AdaptedPackage,
     AdaptOptions,
     AdaptTokenizer,
     DeltaLearnPayload,
@@ -178,7 +181,7 @@ export {
 } from './import/index.js';
 export type { ImportOptions } from './import/index.js';
 
-// ── Foreign weight port (Falcon-Mamba / Codestral-Mamba → HybridMambaModel) ────
+// ── Foreign weight port (Mamba / Falcon-Mamba / Codestral-Mamba → HybridMambaModel) ──
 // Transformer checkpoints are NOT portable — distillation is the only route, and
 // `foreignMambaAdapterFor` rejects them saying so.
 
@@ -191,7 +194,11 @@ export {
     executePortPlan,
     normaliseSourceName,
     falconMambaAdapter,
+    mambaAdapter,
+    mamba1Adapter,
     codestralMambaAdapter,
+    mambaConfigFromTensors,
+    resizePortedVocab,
 } from './import/index.js';
 export type {
     ForeignConfig,
@@ -204,6 +211,7 @@ export type {
     PortedCheckpoint,
     PortedTensors,
     SynthesisedTarget,
+    ResizeVocabOptions,
 } from './import/index.js';
 
 // ── Benchmarking (held-out perplexity / accuracy / throughput + A/B) ──────────
@@ -363,3 +371,8 @@ export type {
 } from './limbic/index.js';
 
 export { LIMBIC_AFFECT_WGSL } from './kernels/limbic_affect.js';
+
+
+// ── Experience (demonstrations, Train Once skills, runs) ─────────────────────
+
+export * from './experience/index.js';

package/src/moe/moe_model.ts

@@ -188,18 +188,20 @@ export class LoadBalanceAccumulator {
     for (const e of route.experts) this.counts[e] = this.counts[e]! + 1;
     for (let e = 0; e < this.numExperts; e++) this.probSum[e] = this.probSum[e]! + route.probs[e]!;
   }
+  /** `f_e`: the fraction of all dispatches that went to each expert (sums to 1 once anything is observed). */
+  dispatchFractions(): Float32Array {
+    const dispatched = this.counts.reduce((a, b) => a + b, 0) || 1; // = tokens·topK
+    return Float32Array.from(this.counts, (c) => c / dispatched);
+  }
   /** The load-balance loss over everything observed so far (0 if no tokens). */
   loss(): number {
     if (this.tokens === 0) return 0;
-    const E = this.numExperts;
-    const dispatched = this.counts.reduce((a, b) => a + b, 0) || 1; // = tokens·topK
+    const f = this.dispatchFractions();
     let sum = 0;
-    for (let e = 0; e < E; e++) {
-      const f = this.counts[e]! / dispatched; // fraction of dispatches to e
-      const p = this.probSum[e]! / this.tokens; // mean router prob for e
-      sum += f * p;
+    for (let e = 0; e < this.numExperts; e++) {
+      sum += f[e]! * (this.probSum[e]! / this.tokens); // f_e · mean router prob for e
     }
-    return E * sum;
+    return this.numExperts * sum;
   }
 }
 

package/src/moe/moe_trainer.ts

@@ -8,7 +8,7 @@
  * WebGPU optimiser kernel would accelerate.
  */
 
-import { SharedExpertMoE } from "./moe_model.js";
+import { LoadBalanceAccumulator, SharedExpertMoE } from "./moe_model.js";
 import { AdamW } from "../optim/adamw.js";
 
 export interface MoESample {
@@ -87,7 +87,7 @@ export class MoETrainer {
       // Forward + task backward, retaining (x, probs) for the batch aux gradient.
       const xs: Float32Array[] = [];
       const probsList: Float32Array[] = [];
-      const counts = new Float32Array(numExperts);
+      const balance = new LoadBalanceAccumulator(numExperts);
       let batchLoss = 0;
 
       for (const s of batch) {
@@ -101,12 +101,11 @@ export class MoETrainer {
         this.model.backward(dOut, f.cache);
         xs.push(f.cache.x);
         probsList.push(f.route.probs);
-        for (const ex of f.route.experts) counts[ex] = counts[ex]! + 1;
+        balance.observe(f.route);
       }
 
       // Load-balancing aux gradient (batch-level): f = dispatch fractions.
-      const dispatched = counts.reduce((a, b) => a + b, 0) || 1;
-      const fVec = Float32Array.from(counts, (c) => c / dispatched);
+      const fVec = balance.dispatchFractions();
       const scale = (this.opt.auxWeight * numExperts) / batch.length;
       for (let i = 0; i < xs.length; i++) {
         this.model.auxGradStep(xs[i]!, probsList[i]!, fVec, scale);
@@ -117,9 +116,7 @@ export class MoETrainer {
       this.adam.step();
 
       epochLoss += batchLoss;
-      lastAux = numExperts * fVec.reduce((sum, f, e) => sum + f * (probsList.length
-        ? probsList.reduce((s, p) => s + p[e]!, 0) / probsList.length
-        : 0), 0);
+      lastAux = balance.loss();
     }
 
     return { loss: epochLoss / Math.max(1, samples.length), auxLoss: lastAux };

package/src/utils/rng.ts

@@ -50,16 +50,20 @@ export function setInitSeed(seed: number | undefined): void {
     }
 }
 
-/** Box–Muller Gaussian sample from the active source. */
-export function randn(std = 1): number {
-    const u1 = Math.max(_next(), 1e-12);
-    const u2 = _next();
+/**
+ * Box–Muller Gaussian sample. Draws from the active init source unless `next` names
+ * another uniform [0,1) source — a caller-owned {@link SeededRng}, say, when the draw
+ * must be reproducible without touching the process-wide seed.
+ */
+export function randn(std = 1, next: () => number = _next): number {
+    const u1 = Math.max(next(), 1e-12);
+    const u2 = next();
     return std * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
 }
 
 /** Returns a Float32Array of `n` Gaussian samples with the given standard deviation. */
-export function gaussianArray(n: number, std: number): Float32Array {
+export function gaussianArray(n: number, std: number, next?: () => number): Float32Array {
     const a = new Float32Array(n);
-    for (let i = 0; i < n; i++) a[i] = randn(std);
+    for (let i = 0; i < n; i++) a[i] = randn(std, next);
     return a;
 }