mlx-bun 0.2.0 → 0.4.0

package/src/weights.ts

@@ -59,12 +59,30 @@ const NATIVE_WEIGHTS: WeightsNativeBindings = {
     ),
 };
 
+/**
+ * Checkpoint-generation normalization, installed once at load.
+ *
+ * `names` maps the GRAPH's canonical tensor name onto the name the artifact
+ * actually stores; a canonical name absent from the map does not exist (the
+ * same way a reference `sanitize()` drops a tensor). `fixup` is an optional
+ * lazy value repair applied to the loaded array — return null to keep it.
+ *
+ * Weights only APPLIES a view; the rules belong to the architecture that owns
+ * the naming (see `qwen35WeightsView` in model/qwen3_5-checkpoint.ts).
+ */
+export interface WeightsView {
+  readonly names: ReadonlyMap<string, string>;
+  fixup?(canonical: string, arr: MlxArray): MlxArray | null;
+}
+
 export class Weights {
   readonly shards: ShardedSafetensors;
   /** shard filename → native mlx map handle (string → lazy array). */
   readonly #maps = new Map<string, bigint>();
+  /** Cached arrays, keyed by CANONICAL name (post-view). */
   readonly #arrays = new Map<string, MlxArray>();
   readonly #native: WeightsNativeBindings;
+  #view: WeightsView | null = null;
 
   private constructor(
     shards: ShardedSafetensors,
@@ -120,6 +138,19 @@ export class Weights {
     }
   }
 
+  /** Install the checkpoint normalization. Must precede the first tensor()
+   *  call: the array cache and every shard lookup below are keyed through it. */
+  setView(view: WeightsView): void {
+    if (this.#arrays.size > 0)
+      throw new Error("Weights.setView after tensors were materialized");
+    this.#view = view;
+  }
+
+  /** Canonical name → the stored name (identity when no view is installed). */
+  #source(name: string): string {
+    return this.#view ? this.#view.names.get(name) ?? name : name;
+  }
+
   /** Free one shard's native map AND every cached tensor from it, releasing
    *  all bytes materialized while consuming that shard. The shard transparently
    *  re-opens on the next tensor() touching it (header parse only — lazy).
@@ -130,7 +161,7 @@ export class Weights {
     const sf = this.shards.files.get(file);
     if (!sf) return;
     for (const [name, arr] of this.#arrays) {
-      if (this.shards.tensorToFile.get(name) === sf) {
+      if (this.shards.tensorToFile.get(this.#source(name)) === sf) {
         arr.dispose();
         this.#arrays.delete(name);
       }
@@ -144,29 +175,36 @@ export class Weights {
 
   /** The shard filename holding `name` (for releaseShard scheduling). */
   fileOf(name: string): string | undefined {
-    const sf = this.shards.tensorToFile.get(name);
+    const sf = this.shards.tensorToFile.get(this.#source(name));
     if (!sf) return undefined;
     for (const [file, cand] of this.shards.files) if (cand === sf) return file;
     return undefined;
   }
 
+  /** Canonical names — with a view installed, exactly the tensors the graph
+   *  can ask for (renamed, drops excluded). */
   get tensorNames(): string[] {
-    return this.shards.tensorNames;
+    return this.#view ? [...this.#view.names.keys()] : this.shards.tensorNames;
   }
 
   info(name: string): TensorInfo {
-    return this.shards.info(name);
+    return this.shards.info(this.#source(name));
   }
 
   has(name: string): boolean {
-    return this.shards.tensorToFile.has(name);
+    return this.#view
+      ? this.#view.names.has(name)
+      : this.shards.tensorToFile.has(name);
   }
 
-  /** Lazy mlx array for a tensor; cached per name. */
+  /** Lazy mlx array for a tensor; cached per canonical name. */
   tensor(name: string): MlxArray {
     let arr = this.#arrays.get(name);
     if (!arr) {
-      const sf = this.shards.tensorToFile.get(name);
+      if (this.#view && !this.#view.names.has(name))
+        throw new Error(`no tensor named ${name}`);
+      const source = this.#source(name);
+      const sf = this.shards.tensorToFile.get(source);
       if (!sf) throw new Error(`no tensor named ${name}`);
       let entry = [...this.#maps.entries()]
         .find(([file]) => this.shards.files.get(file) === sf);
@@ -179,9 +217,17 @@ export class Weights {
       const mapHandle = entry[1];
       const slot = new BigUint64Array([C.mlx_array_new()]);
       const slotPtr = ptr(slot);
-      if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(name))) !== 0)
-        throw new Error(`tensor ${name} missing from native map`);
+      if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(source))) !== 0)
+        throw new Error(`tensor ${source} missing from native map`);
       arr = new MlxArray(read.u64(slotPtr, 0));
+      // Value repair (γ−1 → γ, conv layout): the fixed array is a lazy graph
+      // node over the source, so the raw handle is released immediately — the
+      // node keeps its own reference to the underlying data.
+      const fixed = this.#view?.fixup?.(name, arr) ?? null;
+      if (fixed) {
+        arr.dispose();
+        arr = fixed;
+      }
       this.#arrays.set(name, arr);
     }
     return arr;

package/docs/archive/investigations/curve-runs/bisector-route.json

@@ -1,98 +0,0 @@
-[
-  {
-    "tag": "creative",
-    "pol": "default",
-    "departedPct": 1,
-    "fr": 0.2571428571428571,
-    "rr": 0.34375,
-    "lr": 0,
-    "ratio": 0.748051948051948,
-    "cohMean": -1.0010227537330092,
-    "cohN": 9,
-    "cohOk": 9
-  },
-  {
-    "tag": "creative",
-    "pol": "temp0.8",
-    "departedPct": 1,
-    "fr": 0.06976744186046512,
-    "rr": 0.2318840579710145,
-    "lr": 0.022727272727272728,
-    "ratio": 0.3008720930232558,
-    "cohMean": -0.8618866522327556,
-    "cohN": 3,
-    "cohOk": 3
-  },
-  {
-    "tag": "creative",
-    "pol": "curve-id",
-    "departedPct": 1,
-    "fr": 0.2647058823529412,
-    "rr": 0.34375,
-    "lr": 0,
-    "ratio": 0.7700534759358288,
-    "cohMean": -1.079108765834491,
-    "cohN": 9,
-    "cohOk": 9
-  },
-  {
-    "tag": "creative",
-    "pol": "curve-router",
-    "departedPct": 1,
-    "fr": 0.43333333333333335,
-    "rr": 0.5454545454545454,
-    "lr": 0.03225806451612903,
-    "ratio": 0.7944444444444445,
-    "cohMean": -1.3585284007349332,
-    "cohN": 13,
-    "cohOk": 13
-  },
-  {
-    "tag": "factual",
-    "pol": "default",
-    "departedPct": 1,
-    "fr": 0.42105263157894735,
-    "rr": 0.021164021164021163,
-    "lr": 0,
-    "ratio": 19.894736842105264,
-    "cohMean": -0.5009498243893914,
-    "cohN": 16,
-    "cohOk": 16
-  },
-  {
-    "tag": "factual",
-    "pol": "temp0.8",
-    "departedPct": 1,
-    "fr": 0.38461538461538464,
-    "rr": 0.02564102564102564,
-    "lr": 0,
-    "ratio": 15.000000000000002,
-    "cohMean": -0.42202286759264085,
-    "cohN": 15,
-    "cohOk": 15
-  },
-  {
-    "tag": "factual",
-    "pol": "curve-id",
-    "departedPct": 1,
-    "fr": 0.40540540540540543,
-    "rr": 0.026881720430107527,
-    "lr": 0,
-    "ratio": 15.081081081081082,
-    "cohMean": -0.5881801143434967,
-    "cohN": 15,
-    "cohOk": 15
-  },
-  {
-    "tag": "factual",
-    "pol": "curve-router",
-    "departedPct": 1,
-    "fr": 0.5714285714285714,
-    "rr": 0.0440251572327044,
-    "lr": 0.04950495049504951,
-    "ratio": 12.979591836734693,
-    "cohMean": -0.770856366133446,
-    "cohN": 8,
-    "cohOk": 8
-  }
-]