@genai-fi/nanogpt 0.8.1 → 0.8.3

package/dist/Generator.d.ts

@@ -45,7 +45,11 @@ export default class Generator extends EE<'start' | 'stop' | 'tokens'> {
     generate(prompt?: string, options?: IGenerateOptions): Promise<string>;
     stop(): void;
     getText(): string;
-    getAttentionData(): number[][][][];
+    getAttentionData(): number[][][][][];
     getProbabilitiesData(): number[][][];
+    getEmbeddingsData(): {
+        name: string;
+        tensor: number[][];
+    }[][];
     getTokens(): number[];
 }

package/dist/Generator.js

@@ -63,18 +63,18 @@ import { c as G } from "./concat-pHiVqR3L.js";
  * limitations under the License.
  * =============================================================================
  */
-function N(m, t, e, i = !1) {
-  const o = L(m, "logits", "multinomial"), s = o.size, n = o.rank;
+function N(h, t, e, i = !1) {
+  const o = L(h, "logits", "multinomial"), s = o.size, n = o.rank;
   if (s < 2)
     throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
   if (n > 2)
     throw new Error(`Rank of probabilities must be 1 or 2, but is ${n}`);
   e = e || Math.random();
-  const a = { logits: n === 1 ? x(o, [1, -1]) : o }, p = { numSamples: t, seed: e, normalized: i }, l = C.runKernel(I, a, p);
-  return n === 1 ? x(l, [l.size]) : l;
+  const a = { logits: n === 1 ? x(o, [1, -1]) : o }, l = { numSamples: t, seed: e, normalized: i }, m = C.runKernel(I, a, l);
+  return n === 1 ? x(m, [m.size]) : m;
 }
-const S = /* @__PURE__ */ A({ multinomial_: N }), H = [
-  ...Array.from({ length: 95 }, (m, t) => String.fromCharCode(t + 32)),
+const D = /* @__PURE__ */ A({ multinomial_: N }), H = [
+  ...Array.from({ length: 95 }, (h, t) => String.fromCharCode(t + 32)),
   // ASCII
   // Spanish accented letters and punctuation
   ..."áéíóúüñ¿¡",
@@ -85,8 +85,8 @@ const S = /* @__PURE__ */ A({ multinomial_: N }), H = [
   // Cyrillic letters
   ..."абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ"
 ];
-function U(m, t) {
-  return m.length === t ? m : m.length > t ? m.slice(0, t) : m.concat(Array(t - m.length).fill(""));
+function U(h, t) {
+  return h.length === t ? h : h.length > t ? h.slice(0, t) : h.concat(Array(t - h.length).fill(""));
 }
 class qt extends z {
   constructor(t, e) {
@@ -112,7 +112,9 @@ class qt extends z {
       return null;
     const n = await t.decode([s]);
     if (i) {
-      const d = await Promise.all(i.map((a) => a.array().then((p) => p)));
+      const d = await Promise.all(
+        i.map((a) => a.array().then((l) => l))
+      );
       i.forEach((a) => a.dispose()), this.attentionData.push(d);
     }
     if (o) {
@@ -130,48 +132,47 @@ class qt extends z {
       } : void 0,
       cache: e,
       outputEmbeddings: i?.embeddings ?? !1
-    }, p = O(() => {
-      const r = t, h = r.shape[1], u = h <= this.model.config.blockSize ? r : r.slice(
-        [0, h - this.model.config.blockSize],
-        [r.shape[0], this.model.config.blockSize]
-      ), k = d ? this.model.config.blockSize - u.shape[1] : 0, b = k > 0 ? _(u, [
+    }, l = O(() => {
+      const p = t, u = p.shape[1], r = u <= this.model.config.blockSize ? p : p.slice(
+        [0, u - this.model.config.blockSize],
+        [p.shape[0], this.model.config.blockSize]
+      ), f = d ? this.model.config.blockSize - r.shape[1] : 0, b = f > 0 ? _(r, [
         [0, 0],
-        [0, k]
-      ]) : u, [f] = this.model.forward(a, b), y = f.shape[1] - 1 - k, c = f.slice([0, y, 0], [f.shape[0], 1, f.shape[2]]);
+        [0, f]
+      ]) : r, [g] = this.model.forward(a, b), y = g.shape[1] - 1 - f, c = g.slice([0, y, 0], [g.shape[0], 1, g.shape[2]]);
       return a.attentionScores?.attentionOut && a.attentionScores.attentionOut.forEach((T, E) => {
         T.shape[1] !== 1 && (a.attentionScores.attentionOut[E] = R(
           T.slice([0, y, 0], [T.shape[0], 1, T.shape[2]])
         ), T.dispose());
-      }), f.dispose(), c.div(o).squeeze([1]);
+      }), g.dispose(), c.div(o).squeeze([1]);
     });
-    let l;
+    let m;
     if (n) {
-      const r = v(p), h = await r.array();
-      r.dispose();
-      const u = h[0].map((c, g) => ({ prob: c, index: g })).sort((c, g) => g.prob - c.prob);
-      let k = 0;
-      const b = new Array(u.length).fill(0);
-      for (const c of u)
-        if (k += c.prob, b[c.index] = c.prob, k >= n)
+      const p = v(l), u = await p.array();
+      p.dispose();
+      const r = u[0].map((c, k) => ({ prob: c, index: k })).sort((c, k) => k.prob - c.prob);
+      let f = 0;
+      const b = new Array(r.length).fill(0);
+      for (const c of r)
+        if (f += c.prob, b[c.index] = c.prob, f >= n)
           break;
-      const f = b.reduce((c, g) => c + g, 0), y = b.map((c) => c / f);
-      l = $(y);
+      const g = b.reduce((c, k) => c + k, 0), y = b.map((c) => c / g);
+      m = $(y);
     } else if (s) {
-      const { values: r, indices: h } = K(p, s), u = S(r, 1);
-      l = q(h, u, 1), r.dispose(), h.dispose(), u.dispose();
+      const { values: p, indices: u } = K(l, s), r = D(p, 1);
+      m = q(u, r, 1), p.dispose(), u.dispose(), r.dispose();
     } else
-      l = S(p, 1);
+      m = D(l, 1);
     let w;
-    i?.includeProbabilities && (w = v(p)), a.embeddings && this.embeddingsData.push(
-      await Promise.all(
-        a.embeddings.map(async (r) => {
-          const h = await r.array();
-          return r.dispose(), h;
-        })
-      )
-    );
-    const D = l.reshape([1, 1]);
-    return l.dispose(), l = D, p.dispose(), { output: l, probabilities: w, attention: a.attentionScores?.attentionOut };
+    if (i?.includeProbabilities && (w = v(l)), a.embeddings) {
+      const p = a.embeddings.map(async (r) => {
+        const f = await r.tensor.array();
+        return r.tensor.dispose(), { name: r.name, tensor: f };
+      }), u = await Promise.all(p);
+      this.embeddingsData.push(u);
+    }
+    const S = m.reshape([1, 1]);
+    return m.dispose(), m = S, l.dispose(), { output: m, probabilities: w, attention: a.attentionScores?.attentionOut };
   }
   /** Generate multiple tokens in a loop and produce text */
   async _generate(t) {
@@ -189,8 +190,8 @@ class qt extends z {
       if (this.cache)
         e.dispose(), e = s;
       else {
-        const p = e;
-        e = G([e, s], 1), p.dispose();
+        const l = e;
+        e = G([e, s], 1), l.dispose();
       }
       const a = await this.processResponse(this.actualTokeniser, s, d, n);
       if (this.cache || s.dispose(), a === null)
@@ -239,6 +240,9 @@ class qt extends z {
   getProbabilitiesData() {
     return this.probabilitiesData;
   }
+  getEmbeddingsData() {
+    return this.embeddingsData;
+  }
   getTokens() {
     return this.tokens;
   }

package/dist/checks/index.d.ts

@@ -6,6 +6,7 @@ import { execute as normRMSGrad } from './normRMSGrad';
 import { execute as appendCache } from './appendCache';
 import { execute as attentionMask } from './attentionMask';
 import { default as runCheck } from './check';
+import { createWeightStatistics, createTensorStatistics } from './weights';
 declare const checks: {
     rope: typeof rope;
     qkv: typeof qkv;
@@ -15,5 +16,7 @@ declare const checks: {
     appendCache: typeof appendCache;
     attentionMask: typeof attentionMask;
     runCheck: typeof runCheck;
+    createLayerWeightStatistics: typeof createWeightStatistics;
+    createWeightStatistics: typeof createTensorStatistics;
 };
 export default checks;

package/dist/checks/index.js

@@ -1,20 +1,23 @@
 import { execute as e } from "./rope.js";
 import { execute as t } from "./normRMS.js";
-import { execute as o } from "./qkv.js";
-import { execute as r } from "./gelu.js";
-import { execute as c } from "./normRMSGrad.js";
-import { execute as m } from "./appendCache.js";
-import { execute as u } from "./attentionMask.js";
-import x from "./check.js";
+import { execute as r } from "./qkv.js";
+import { execute as c } from "./gelu.js";
+import { execute as o } from "./normRMSGrad.js";
+import { execute as a } from "./appendCache.js";
+import { execute as i } from "./attentionMask.js";
+import m from "./check.js";
+import { createTensorStatistics as s, createWeightStatistics as u } from "./weights.js";
 const d = {
   rope: e,
-  qkv: o,
-  gelu: r,
+  qkv: r,
+  gelu: c,
   normRMS: t,
-  normRMSGrad: c,
-  appendCache: m,
-  attentionMask: u,
-  runCheck: x
+  normRMSGrad: o,
+  appendCache: a,
+  attentionMask: i,
+  runCheck: m,
+  createLayerWeightStatistics: u,
+  createWeightStatistics: s
 };
 export {
   d as default

package/dist/checks/qkv.js

@@ -1,9 +1,34 @@
-import { s as c, e as m } from "../index-DdmHGZjq.js";
-import { t as i } from "../tensor3d-BOukqWwr.js";
-import { t as u } from "../tensor2d-CObBWBkW.js";
-async function w(a) {
-  await c(a);
-  const o = i(
+import { z as i, A as u, B as c, s as l, e as h } from "../index-DdmHGZjq.js";
+import { t as f } from "../tensor2d-CObBWBkW.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function m(t, e, n) {
+  if (i(t), e != null && e.length !== 3)
+    throw new Error("tensor3d() requires shape to have three numbers");
+  const r = u(t, n);
+  if (r.length !== 3 && r.length !== 1)
+    throw new Error("tensor3d() requires values to be number[][][] or flat/TypedArray");
+  if (r.length === 1 && e == null)
+    throw new Error("tensor3d() requires shape to be provided when `values` are a flat array");
+  return c(t, e, r, n);
+}
+async function y(t) {
+  await l(t);
+  const e = m(
     [
       [
         [0.1, 0.2],
@@ -11,15 +36,15 @@ async function w(a) {
       ]
     ],
     [1, 2, 2]
-  ), r = u(
+  ), n = f(
     [
       [0.5, 0.6, 0.9, 1, 1.3, 1.4],
       [0.7, 0.8, 1.1, 1.2, 1.5, 1.6]
     ],
     [2, 6]
-  ), t = m().runKernel("QKV", { x: o, kernel: r }, { heads: 1 }), s = await t[0].array(), n = await t[1].array(), e = await t[2].array();
-  return [s, n, e];
+  ), r = h().runKernel("QKV", { x: e, kernel: n }, { heads: 1 }), o = await r[0].array(), a = await r[1].array(), s = await r[2].array();
+  return [o, a, s];
 }
 export {
-  w as execute
+  y as execute
 };

package/dist/checks/rope.js

@@ -1,9 +1,9 @@
+import t from "../layers/RoPECache.js";
 import { s as c, e as i } from "../index-DdmHGZjq.js";
-import { t as m } from "../tensor4d-DLtk7Nxh.js";
-import { t } from "../tensor3d-BOukqWwr.js";
-async function y(n) {
-  await c(n);
-  const s = m(
+import { t as p } from "../tensor4d-DLtk7Nxh.js";
+async function y(a) {
+  await c(a);
+  const o = p(
     [
       [
         [
@@ -13,8 +13,25 @@ async function y(n) {
       ]
     ],
     [1, 1, 2, 2]
-  ), e = t([0.5, 0.6], [2, 1, 1]), o = t([0.9, 1], [2, 1, 1]), r = i().runKernel("Rope", { x: s, sin: e, cos: o }, { pastLen: 0 });
-  return Array.isArray(r) ? r.map((a) => a.array()) : r.array();
+  ), n = {
+    biasInLayerNorm: !1,
+    vocabSize: 20,
+    nEmbed: 16,
+    nHead: 2,
+    nLayer: 1,
+    biasInLinear: !1,
+    dropout: 0,
+    blockSize: 128,
+    mlpFactor: 4,
+    useRope: !0
+  }, e = new t(n);
+  e.ensureRopeCache(120);
+  const r = i().runKernel(
+    "Rope",
+    { x: o, sin: e.getSin(), cos: e.getCos() },
+    { pastLen: 20 }
+  );
+  return Array.isArray(r) ? r.map((s) => s.array()) : r.array();
 }
 export {
   y as execute

package/dist/checks/weights.d.ts

@@ -0,0 +1,16 @@
+import { default as BaseLayer } from '../layers/BaseLayer';
+import { Tensor } from '@tensorflow/tfjs-core';
+export interface TensorStatistics {
+    mean: number;
+    std: number;
+    min: number;
+    max: number;
+    sparsity: number;
+    isFinite: boolean;
+    hasNaN: boolean;
+    closeToZeroCount: number;
+}
+export declare function createTensorStatistics(weight: Tensor | number[]): Promise<TensorStatistics>;
+export declare function createWeightStatistics(layer: BaseLayer): Promise<{
+    [key: string]: TensorStatistics;
+}>;

package/dist/checks/weights.js

@@ -0,0 +1,29 @@
+async function d(s) {
+  const e = Array.isArray(s) ? s : await s.data(), a = e.length;
+  let n = 0, l = 0, i = e[0], r = e[0], u = 0, h = !0, f = !1, m = 0;
+  for (let c = 0; c < a; c++) {
+    const t = e[c];
+    n += t, l += t * t, t < i && (i = t), t > r && (r = t), t === 0 && u++, Math.abs(t) < 1e-8 && m++, Number.isNaN(t) && (f = !0);
+  }
+  const o = n / a, y = l / a - o * o, N = Math.sqrt(y), b = u / a;
+  return {
+    mean: o,
+    std: N,
+    min: i,
+    max: r,
+    sparsity: b,
+    isFinite: h,
+    hasNaN: f,
+    closeToZeroCount: m
+  };
+}
+async function S(s) {
+  const e = s.trainableVariables, a = {};
+  for (const n of e)
+    a[n.name] = await d(n);
+  return a;
+}
+export {
+  d as createTensorStatistics,
+  S as createWeightStatistics
+};

package/dist/layers/BaseLayer.d.ts

@@ -6,6 +6,11 @@ export interface ForwardAttributes {
     training: boolean;
     checkpointing?: boolean;
     ropeCache?: RoPECache;
+    outputEmbeddings?: boolean;
+    embeddings?: {
+        name: string;
+        tensor: Tensor;
+    }[];
 }
 export default abstract class BaseLayer<ATTR extends ForwardAttributes = ForwardAttributes> {
     readonly parent?: BaseLayer;

package/dist/layers/CausalSelfAttention.js

@@ -1,16 +1,16 @@
 import { attentionMask as g } from "../ops/attentionMask.js";
 import O from "./BaseLayer.js";
-import { qkv as P } from "../ops/qkv.js";
+import { qkv as _ } from "../ops/qkv.js";
 import { rope as v } from "../ops/rope.js";
 import { appendCache as V } from "../ops/appendCache.js";
-import { k as c, t as C } from "../index-DdmHGZjq.js";
-import { fusedSoftmax as T } from "../ops/fusedSoftmax.js";
-import { d as L } from "../random_width-DKGeiFuR.js";
-import { v as b } from "../variable-DPFOJyRG.js";
-import { r as k, d as y } from "../dropout-CcKSfOYE.js";
-import { r as N } from "../reshape-WeJkT3ja.js";
-import { m as R } from "../mat_mul-Dpy2mMRu.js";
-class $ extends O {
+import { k as c, t as P } from "../index-DdmHGZjq.js";
+import { fusedSoftmax as b } from "../ops/fusedSoftmax.js";
+import { d as C } from "../random_width-DKGeiFuR.js";
+import { v as k } from "../variable-DPFOJyRG.js";
+import { r as T, d as L } from "../dropout-CcKSfOYE.js";
+import { r as j } from "../reshape-WeJkT3ja.js";
+import { m as x } from "../mat_mul-Dpy2mMRu.js";
+class W extends O {
   divisor;
   index;
   units;
@@ -23,34 +23,34 @@ class $ extends O {
   build() {
     this.hasVariable(this.ATTN) === !1 && this.setVariable(
       this.ATTN,
-      b(
-        k([this.config.nEmbed, this.units], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cAttn_kernel`
+      k(
+        T([this.config.nEmbed, this.units], 0, 0.02),
+        !0,
+        `block_${this.index}_attn_cAttn_kernel`
       )
     ), this.hasVariable(this.PROJ) === !1 && this.setVariable(
       this.PROJ,
-      b(
-        k([this.projUnits, this.config.nEmbed], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cProj_kernel`
+      k(
+        T([this.projUnits, this.config.nEmbed], 0, 0.02),
+        !0,
+        `block_${this.index}_attn_cProj_kernel`
       )
     );
   }
   getAttentionScores(t, i, s, o) {
-    const e = g(t, i, this.divisor), n = T(e, s ? this.config.dropout : 0, o);
+    const e = g(t, i, this.divisor), n = b(e, s ? this.config.dropout : 0, o);
     return e.dispose(), n;
   }
   // Attention with optional past. If pastLen > 0 and T_cur == 1, no mask needed.
   getAttentionScoresWithPast(t, i, s) {
-    const o = g(t, i, this.divisor, s), e = T(o, 0, 0);
+    const o = g(t, i, this.divisor, s), e = b(o, 0, 0);
     return o.dispose(), e;
   }
   getQKV(t) {
-    return P(t, this.getVariable(this.ATTN), this.config.nHead);
+    return _(t, this.getVariable(this.ATTN), this.config.nHead);
   }
   getOutputProjection(t) {
-    const i = t.shape[0], s = t.shape[2], o = this.config.nEmbed, e = t.transpose([0, 2, 1, 3]), n = N(e, [i, s, o]), r = L(n, this.getVariable(this.PROJ));
+    const i = t.shape[0], s = t.shape[2], o = this.config.nEmbed, e = t.transpose([0, 2, 1, 3]), n = j(e, [i, s, o]), r = C(n, this.getVariable(this.PROJ));
     return n.dispose(), e.dispose(), r;
   }
   updateCache(t, i, s) {
@@ -62,19 +62,19 @@ class $ extends O {
     s.length = d, s.cumulativeLength = h, s.k = c(r), s.v = c(p);
   }
   forward(t, i) {
-    return C(() => {
+    return P(() => {
       this.startMemory();
       const [s, o, e] = this.getQKV(i), n = t.pastKV ? t.pastKV.cumulativeLength : 0, r = t.ropeCache, p = r ? v(s, r, n) : s, d = r ? v(o, r, n) : o;
       r && (s.dispose(), o.dispose());
       const h = t.pastKV ? t.pastKV.length : 0;
       t.pastKV && !t.training && this.updateCache(d, e, t.pastKV);
-      const u = t.pastKV?.k ? t.pastKV.k : d, m = t.pastKV?.v ? t.pastKV.v : e;
+      const u = t.pastKV?.k ? t.pastKV.k : d, l = t.pastKV?.v ? t.pastKV.v : e;
       let a;
       h > 0 ? a = this.getAttentionScoresWithPast(p, u, h) : a = this.getAttentionScores(p, u, t.training, t.seed || 0), p.dispose(), t.pastKV || u.dispose();
-      const l = R(a, m), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
-      f || a.dispose(), t.pastKV || m.dispose();
-      const A = this.getOutputProjection(l);
-      if (l.dispose(), f && t.attentionScores && t.attentionScores.attentionOut !== void 0) {
+      const m = x(a, l), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
+      f || a.dispose(), t.pastKV || l.dispose();
+      const A = this.getOutputProjection(m);
+      if (m.dispose(), f && t.attentionScores && t.attentionScores.attentionOut !== void 0) {
         const K = a.shape[1], S = a.shape[2];
         t.attentionScores.attentionOut?.push(
           c(a.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
@@ -85,12 +85,12 @@ class $ extends O {
   }
   dropout(t) {
     if (this.config.dropout > 0) {
-      const i = y(t, this.config.dropout);
+      const i = L(t, this.config.dropout);
       return t.dispose(), i;
     } else
       return t;
   }
 }
 export {
-  $ as default
+  W as default
 };

package/dist/layers/MLP.js

@@ -1,10 +1,10 @@
 import { t as p } from "../index-DdmHGZjq.js";
 import u from "./BaseLayer.js";
 import { matMulGelu as M } from "../ops/matMulGelu.js";
-import { v as o } from "../variable-DPFOJyRG.js";
-import { r as h, d as f } from "../dropout-CcKSfOYE.js";
-import { r as d } from "../reshape-WeJkT3ja.js";
-import { m as c } from "../mat_mul-Dpy2mMRu.js";
+import { v as a } from "../variable-DPFOJyRG.js";
+import { r as d, d as c } from "../dropout-CcKSfOYE.js";
+import { r as h } from "../reshape-WeJkT3ja.js";
+import { m as b } from "../mat_mul-Dpy2mMRu.js";
 class H extends u {
   index;
   hiddenUnits;
@@ -16,32 +16,32 @@ class H extends u {
   build() {
     this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(
       this.MLPHIDDEN,
-      o(
-        h([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cAttn_kernel`
+      a(
+        d([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
+        !0,
+        `block_${this.index}_mlpHidden_kernel`
       )
     ), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(
       this.MLPOUT,
-      o(
-        h([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
-        !0
-        //`block_${this.index}_attn_cProj_kernel`
+      a(
+        d([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
+        !0,
+        `block_${this.index}_mlpOut_kernel`
       )
     );
   }
   forward(i, t) {
     return p(() => {
       this.startMemory();
-      const [s, r, e] = t.shape, n = d(t, [s * r, e]), a = M(n, this.getVariable(this.MLPHIDDEN)), m = c(a, this.getVariable(this.MLPOUT));
-      a.dispose();
-      const l = d(m, [s, r, e]);
-      return this.endMemory("MLP"), l;
+      const [s, e, r] = t.shape, n = h(t, [s * e, r]), o = M(n, this.getVariable(this.MLPHIDDEN)), l = b(o, this.getVariable(this.MLPOUT));
+      o.dispose();
+      const m = h(l, [s, e, r]);
+      return this.endMemory("MLP"), m;
     });
   }
   dropout(i) {
     if (this.config.dropout > 0) {
-      const t = f(i, this.config.dropout);
+      const t = c(i, this.config.dropout);
       return i.dispose(), t;
     }
     return i;

package/dist/layers/TiedEmbedding.js

@@ -2,28 +2,28 @@ import { d as r } from "../random_width-DKGeiFuR.js";
 import "../index-DdmHGZjq.js";
 import { r as a } from "../exports_initializers-DKk7-bsx.js";
 import s from "./BaseLayer.js";
-import { v as m } from "../variable-DPFOJyRG.js";
-import { g as o } from "../gather-CPg6ZlQA.js";
+import { v as o } from "../variable-DPFOJyRG.js";
+import { g as m } from "../gather-CPg6ZlQA.js";
 class S extends s {
   vocabSize;
   embedDim;
   initializer;
   WEIGHTS;
-  constructor(i, e, t) {
-    super(i, t), this.WEIGHTS = e, this.vocabSize = i.vocabSize, this.embedDim = i.nEmbed, this.initializer = a({
+  constructor(i, t, e) {
+    super(i, e), this.WEIGHTS = t, this.vocabSize = i.vocabSize, this.embedDim = i.nEmbed, this.initializer = a({
       mean: 0,
       stddev: 0.02
-    }), this.addVariable(this.WEIGHTS, m(this.initializer.apply([this.vocabSize, this.embedDim]), !0));
+    }), this.addVariable(this.WEIGHTS, o(this.initializer.apply([this.vocabSize, this.embedDim]), !0, t));
   }
   embed(i) {
-    return o(this.getVariable(this.WEIGHTS), i, 0);
+    return m(this.getVariable(this.WEIGHTS), i, 0);
   }
   project(i) {
     return r(i, this.getVariable(this.WEIGHTS).transpose());
   }
   // Dummy, should not be used.
-  forward(i, e) {
-    return this.project(e);
+  forward(i, t) {
+    return this.project(t);
   }
 }
 export {

package/dist/layers/TransformerBlock.js

@@ -1,32 +1,32 @@
 import l from "./CausalSelfAttention.js";
-import r from "./MLP.js";
+import p from "./MLP.js";
 import o from "./RMSNorm.js";
-import d from "./BaseLayer.js";
-import { t as p } from "../index-DdmHGZjq.js";
-class k extends d {
+import m from "./BaseLayer.js";
+import { k as n, t as h } from "../index-DdmHGZjq.js";
+class k extends m {
   ln1;
   attn;
   ln2;
   mlp;
   index;
   skipped = !1;
-  constructor(t, s, i) {
-    super(s, i), this.index = t, this.ln1 = new o(s, `block_${this.index}_rms1`, this), this.attn = new l(this.index, s, this), this.ln2 = new o(s, `block_${this.index}_rms2`, this), this.mlp = new r(this.index, s, this);
+  constructor(i, s, e) {
+    super(s, e), this.index = i, this.ln1 = new o(s, `block_${this.index}_rms1`, this), this.attn = new l(this.index, s, this), this.ln2 = new o(s, `block_${this.index}_rms2`, this), this.mlp = new p(this.index, s, this);
   }
-  getMLPOutput(t, s) {
-    const i = this.ln2.call({ training: s }, t), e = this.mlp.call({ training: s }, i);
-    i.dispose();
-    const n = t.add(e);
-    return t.dispose(), e.dispose(), n;
+  getMLPOutput(i, s) {
+    const e = this.ln2.call({ training: s.training }, i), t = this.mlp.call({ training: s.training }, e);
+    s.outputEmbeddings ? (n(e), s.embeddings.push({ name: `block_ln2_${this.index}`, tensor: e })) : e.dispose();
+    const d = i.add(t);
+    return i.dispose(), s.outputEmbeddings ? (n(t), s.embeddings.push({ name: `block_mlp_out_${this.index}`, tensor: t })) : t.dispose(), d;
   }
-  forward(t, s) {
-    return p(() => {
+  forward(i, s) {
+    return h(() => {
       if (this.skipped)
         return s;
-      const i = this.ln1.call(t, s), e = this.attn.call(t, i);
-      i.dispose();
-      const n = s.add(e);
-      return e.dispose(), this.getMLPOutput(n, t.training);
+      const e = this.ln1.call(i, s), t = this.attn.call(i, e);
+      i.outputEmbeddings ? (n(e), i.embeddings.push({ name: `block_ln1_${this.index}`, tensor: e })) : e.dispose();
+      const d = s.add(t);
+      return i.outputEmbeddings ? (n(t), i.embeddings.push({ name: `block_attn_out_${this.index}`, tensor: t })) : t.dispose(), this.getMLPOutput(d, i);
     });
   }
   dispose() {

package/dist/main.d.ts

@@ -26,3 +26,4 @@ export declare const layers: {
 };
 export { default as AdamExt } from './training/AdamExt';
 export { default as checks } from './checks';
+export type { TensorStatistics } from './checks/weights';

package/dist/models/NanoGPTV1.js

@@ -1,6 +1,6 @@
-import { defaultConfig as m } from "./config.js";
-import f from "../layers/TransformerBlock.js";
-import u from "../layers/TiedEmbedding.js";
+import { defaultConfig as a } from "./config.js";
+import u from "../layers/TransformerBlock.js";
+import f from "../layers/TiedEmbedding.js";
 import g from "../layers/RoPECache.js";
 import b from "../layers/RMSNorm.js";
 import { t as l, k as p } from "../index-DdmHGZjq.js";
@@ -17,9 +17,9 @@ class R extends w {
   // Final layer norm
   ropeCache;
   constructor(e = {}) {
-    super({ ...m, ...e }), this.wte = new u(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new k(this.config, "positional_embedding", this) : this.ropeCache = new g(this.config), this.blocks = [];
+    super({ ...a, ...e }), this.wte = new f(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new k(this.config, "positional_embedding", this) : this.ropeCache = new g(this.config), this.blocks = [];
     for (let i = 0; i < this.config.nLayer; i++)
-      this.blocks.push(new f(i, this.config, this));
+      this.blocks.push(new u(i, this.config, this));
     this.lnF = new b(this.config, "final_rms_norm", this);
   }
   getClassName() {
@@ -47,17 +47,15 @@ class R extends w {
         );
       for (let t = 0; t < this.blocks.length; t++) {
         const c = this.blocks[t], d = Math.random() * 1e9, r = {
-          ropeCache: e.ropeCache,
-          training: e.training,
+          ...e,
           seed: d,
-          attentionScores: e.attentionScores,
           pastKV: e.cache ? e.cache[t] : void 0
-        }, a = e.checkpointing && e.training ? c.callCheckpoint(r, o) : c.call(r, o);
-        e.outputEmbeddings ? (p(o), e.embeddings.push(o)) : o.dispose(), o = a;
+        }, m = e.checkpointing && e.training ? c.callCheckpoint(r, o) : c.call(r, o);
+        e.outputEmbeddings ? (p(o), e.embeddings.push({ name: `block_output_${t}`, tensor: o })) : o.dispose(), o = m;
       }
       o = this.lnF.call(e, o);
       const n = this.wte.project(o);
-      e.outputEmbeddings ? (p(o), e.embeddings.push(o)) : o.dispose();
+      e.outputEmbeddings ? (p(o), e.embeddings.push({ name: "final_norm_output", tensor: o })) : o.dispose();
       let h;
       return s && (h = this.calculateLoss(n, s)), this.endMemory("Forward"), h ? [n, h] : [n];
     });

package/dist/models/model.d.ts

@@ -5,8 +5,6 @@ export interface ModelForwardAttributes extends ForwardAttributes {
     cache?: KVCache[];
     attentionScores?: AttentionScores;
     seed?: number;
-    outputEmbeddings?: boolean;
-    embeddings?: Tensor[];
 }
 interface TrainingState {
     steps: number;

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.8.1",
+    "version": "0.8.3",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",

package/dist/tensor3d-BOukqWwr.js

@@ -1,30 +0,0 @@
-import { z as o, A as a, B as s } from "./index-DdmHGZjq.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
-function h(n, r, t) {
-  if (o(n), r != null && r.length !== 3)
-    throw new Error("tensor3d() requires shape to have three numbers");
-  const e = a(n, t);
-  if (e.length !== 3 && e.length !== 1)
-    throw new Error("tensor3d() requires values to be number[][][] or flat/TypedArray");
-  if (e.length === 1 && r == null)
-    throw new Error("tensor3d() requires shape to be provided when `values` are a flat array");
-  return s(n, r, e, t);
-}
-export {
-  h as t
-};