@genai-fi/nanogpt 0.5.1 → 0.5.2

package/dist/Generator.js

@@ -1,72 +1,121 @@
 import { E as u } from "./index-Dwqa6Zy2.js";
-import "./index-iNhkcAEQ.js";
-import { t as d } from "./tensor2d-tSxWdFMH.js";
-import { c as p } from "./concat-Cxbo2sOz.js";
-class w extends u {
-  constructor(i, t) {
-    super(), this.model = i, this.tokeniser = t;
+import "./index-CnHyhpKc.js";
+import "./ops/cpu/attentionMask.js";
+import "./ops/webgl/attentionMask.js";
+import "./ops/grads/attentionMask.js";
+import "./ops/cpu/qkv.js";
+import "./ops/webgl/qkv.js";
+import "./ops/grads/qkv.js";
+import "@tensorflow/tfjs";
+import "./ops/cpu/rope.js";
+import "./ops/webgl/rope.js";
+import "./ops/grads/rope.js";
+import "./ops/cpu/appendCache.js";
+import "./ops/webgl/appendCache.js";
+import "./ops/cpu/fusedSoftmax.js";
+import "./ops/webgl/fusedSoftmax.js";
+import "./ops/grads/fusedSoftmax.js";
+import "./ops/cpu/matMulGelu.js";
+import "./ops/webgl/matMulGelu.js";
+import "./ops/grads/matMulGelu.js";
+import "./ops/cpu/normRMS.js";
+import "./ops/webgl/normRMS.js";
+import "./ops/grads/normRMS.js";
+import "./random_width-DI2h9CMs.js";
+import "./ops/cpu/gatherSub.js";
+import "./ops/webgl/gatherSub.js";
+import "./ops/cpu/scatterSub.js";
+import "./ops/webgl/scatterSub.js";
+import "./jszip.min-CjP2V1VV.js";
+import f from "./tokeniser/CharTokeniser.js";
+import "./dataset-ZHEPJmED.js";
+import "./index-Tf7vU29b.js";
+import "./papaparse.min-C8l2Kvo1.js";
+import "./ops/cpu/gelu.js";
+import "./ops/webgl/gelu.js";
+import "./ops/grads/gelu.js";
+import { t as d } from "./tensor2d-CqtBzOKq.js";
+import { c as g } from "./concat-BRRtq4S2.js";
+const k = [
+  ...Array.from({ length: 95 }, (a, t) => String.fromCharCode(t + 32)),
+  // ASCII
+  // Spanish accented letters and punctuation
+  ..."áéíóúüñ¿¡",
+  // Finnish accented letters
+  ..."äöÄÖÅå",
+  // Greek letters
+  ..."αβγδεζηθικλμνξοπρστυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ",
+  // Cyrillic letters
+  ..."абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ"
+];
+function T(a, t) {
+  return a.length === t ? a : a.length > t ? a.slice(0, t) : a.concat(Array(t - a.length).fill(""));
+}
+class ot extends u {
+  constructor(t, o) {
+    super(), this.model = t, this.tokeniser = o;
   }
   active = !1;
-  async tokenisePrompt(i) {
-    const t = i ? await this.tokeniser.tokenise([i], !0) : [[this.tokeniser.eosToken]];
-    return d(t, [1, t[0].length], "int32");
+  async tokenisePrompt(t, o) {
+    const r = o ? await t.tokenise([o], !0) : [[t.eosToken]];
+    return d(r, [1, r[0].length], "int32");
   }
-  async generateNoCache(i, t) {
-    let s = await this.tokenisePrompt(i), o = i || "";
-    const n = t?.maxLength ?? 1e3;
-    for (let r = 0; r < n && this.active; r++) {
+  async generateNoCache(t, o, r) {
+    let i = await this.tokenisePrompt(t, o), s = o || "";
+    const n = r?.maxLength ?? 1e3;
+    for (let m = 0; m < n && this.active; m++) {
       const {
         output: e,
-        attention: a,
+        attention: p,
         probabilities: c
-      } = this.model.generate(s, void 0, t), h = s;
-      s = p([s, e], 1), h.dispose();
-      const l = await this.processResponse(e, a, c);
+      } = this.model.generate(i, void 0, r), h = i;
+      i = g([i, e], 1), h.dispose();
+      const l = await this.processResponse(t, e, p, c);
       if (e.dispose(), l === null)
         break;
-      o += l;
+      s += l;
     }
-    return s.dispose(), o;
+    return i.dispose(), s;
   }
-  async processResponse(i, t, s) {
-    const o = (await i.array())[0][0];
-    if (o === this.tokeniser.eosToken)
+  async processResponse(t, o, r, i) {
+    const s = (await o.array())[0][0];
+    if (s === this.tokeniser.eosToken)
       return null;
-    const n = await this.tokeniser.decode([o]);
-    let r;
-    t && (r = await Promise.all(t.map((a) => a.array().then((c) => c))), t.forEach((a) => a.dispose()));
+    const n = await t.decode([s]);
+    let m;
+    r && (m = await Promise.all(r.map((p) => p.array().then((c) => c))), r.forEach((p) => p.dispose()));
     let e;
-    return s && (e = await s.array(), s.dispose()), this.emit("tokens", [o], n, r, e), n;
+    return i && (e = await i.array(), i.dispose()), this.emit("tokens", [s], n, m, e), n;
   }
-  async generateCache(i, t) {
-    let s = await this.tokenisePrompt(i), o = i || "";
+  async generateCache(t, o, r) {
+    let i = await this.tokenisePrompt(t, o), s = o || "";
     const n = new Array(this.model.config.gpt.nLayer);
     for (let e = 0; e < this.model.config.gpt.nLayer; e++)
       n[e] = { k: void 0, v: void 0, length: 0, cumulativeLength: 0 };
-    const r = t?.maxLength ?? 1e3;
-    for (let e = 0; e < r && this.active; e++) {
+    const m = r?.maxLength ?? 1e3;
+    for (let e = 0; e < m && this.active; e++) {
       const {
-        output: a,
+        output: p,
         probabilities: c,
         attention: h
-      } = this.model.generate(s, n, {
-        ...t,
+      } = this.model.generate(i, n, {
+        ...r,
         usePadding: !1
       });
-      s.dispose(), s = a;
-      const l = await this.processResponse(a, h, c);
+      i.dispose(), i = p;
+      const l = await this.processResponse(t, p, h, c);
       if (l === null)
         break;
-      o += l;
+      s += l;
     }
     return n.forEach((e) => {
       e && (e.k && e.k.dispose(), e.v && e.v.dispose());
-    }), s.dispose(), o;
+    }), i.dispose(), s;
   }
-  async generate(i, t) {
-    const s = i && i.length > this.model.config.gpt.blockSize ? i.slice(-this.model.config.gpt.blockSize) : i;
+  async generate(t, o) {
+    const r = t && t.length > this.model.config.gpt.blockSize ? t.slice(-this.model.config.gpt.blockSize) : t;
     this.active = !0, this.emit("start");
-    const n = await (this.model.config.gpt.useRope && !t?.noCache ? this.generateCache(s, t) : this.generateNoCache(s, t));
+    const i = this.tokeniser.trained ? this.tokeniser : new f(T(k, this.tokeniser.vocabSize)), n = await (this.model.config.gpt.useRope && !o?.noCache ? this.generateCache(i, r, o) : this.generateNoCache(i, r, o));
     return this.active = !1, this.emit("stop"), n;
   }
   stop() {
@@ -74,5 +123,5 @@ class w extends u {
   }
 }
 export {
-  w as default
+  ot as default
 };

package/dist/NanoGPTModel.d.ts

@@ -13,6 +13,7 @@ export interface TrainingLogEntry {
 export interface GenerateOptions {
     temperature?: number;
     topK?: number;
+    topP?: number;
     usePadding?: boolean;
     attentionScores?: boolean;
     includeProbabilities?: boolean;

package/dist/NanoGPTModel.js

@@ -1,16 +1,18 @@
-import { defaultConfig as L } from "./config.js";
-import v from "./layers/TransformerBlock.js";
-import { E as T, D as q, T as K, r as P, p as _ } from "./TiedEmbedding-DsDRvLB0.js";
-import F from "./layers/RoPECache.js";
-import D from "./layers/RMSNorm.js";
-import { estimateParameterCount as O } from "./utilities/parameters.js";
-import { createSoftmaxCrossEntropyWithGrad as N } from "./training/sparseCrossEntropy.js";
-import { B as R } from "./BaseLayer-BhrMN8JO.js";
-import { o as E, i as d, q as B, E as y, aa as G, ab as V, ac as j, t as w, a9 as A, f as z, F as W } from "./index-iNhkcAEQ.js";
-import { r as C } from "./reshape-DxTPgnwL.js";
-import { r as H } from "./range-BsFU-SNG.js";
-import { g as J } from "./gather-Bxe1Qip8.js";
-import { s as Q } from "./softmax-BjsptB07.js";
+import { defaultConfig as F } from "./config.js";
+import O from "./layers/TransformerBlock.js";
+import { T as N, r as R } from "./TiedEmbedding-DORsPlNL.js";
+import A from "./layers/RoPECache.js";
+import G from "./layers/RMSNorm.js";
+import { estimateParameterCount as j } from "./utilities/parameters.js";
+import { createSoftmaxCrossEntropyWithGrad as B } from "./training/sparseCrossEntropy.js";
+import V from "./layers/BaseLayer.js";
+import { E as H, D as W, p as J } from "./random_width-DI2h9CMs.js";
+import { o as x, j as y, u as Q, E as I, aa as U, ab as X, ac as Y, t as z, a9 as Z, f as L, H as tt } from "./index-CnHyhpKc.js";
+import { r as T } from "./reshape-CTIbqjwm.js";
+import { r as et } from "./range-CkOJ7090.js";
+import { s as q } from "./softmax-DX6qXAbm.js";
+import { t as ot } from "./ops-DzQTmLIl.js";
+import { g as st } from "./gather-BWyutxwi.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -27,13 +29,13 @@ import { s as Q } from "./softmax-BjsptB07.js";
  * limitations under the License.
  * =============================================================================
  */
-function U(h, t) {
-  let e = d(h, "a", "mod"), o = d(t, "b", "mod");
-  [e, o] = B(e, o);
+function nt(l, t) {
+  let e = y(l, "a", "mod"), o = y(t, "b", "mod");
+  [e, o] = Q(e, o);
   const n = { a: e, b: o };
-  return y.runKernel(G, n);
+  return I.runKernel(U, n);
 }
-const X = /* @__PURE__ */ E({ mod_: U });
+const it = /* @__PURE__ */ x({ mod_: nt });
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -50,17 +52,17 @@ const X = /* @__PURE__ */ E({ mod_: U });
  * limitations under the License.
  * =============================================================================
  */
-function Y(h, t, e, o = !1) {
-  const n = d(h, "logits", "multinomial"), s = n.size, i = n.rank;
+function rt(l, t, e, o = !1) {
+  const n = y(l, "logits", "multinomial"), s = n.size, i = n.rank;
   if (s < 2)
     throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
   if (i > 2)
     throw new Error(`Rank of probabilities must be 1 or 2, but is ${i}`);
   e = e || Math.random();
-  const c = { logits: i === 1 ? C(n, [1, -1]) : n }, l = { numSamples: t, seed: e, normalized: o }, a = y.runKernel(V, c, l);
-  return i === 1 ? C(a, [a.size]) : a;
+  const r = { logits: i === 1 ? T(n, [1, -1]) : n }, u = { numSamples: t, seed: e, normalized: o }, c = I.runKernel(X, r, u);
+  return i === 1 ? T(c, [c.size]) : c;
 }
-const I = /* @__PURE__ */ E({ multinomial_: Y });
+const C = /* @__PURE__ */ x({ multinomial_: rt });
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -77,8 +79,8 @@ const I = /* @__PURE__ */ E({ multinomial_: Y });
  * limitations under the License.
  * =============================================================================
  */
-function Z(h, t = 1, e = !0) {
-  const o = d(h, "x", "topk");
+function ct(l, t = 1, e = !0) {
+  const o = y(l, "x", "topk");
   if (o.rank === 0)
     throw new Error("topk() expects the input to be of rank 1 or higher");
   const n = o.shape[o.shape.length - 1];
@@ -86,10 +88,10 @@ function Z(h, t = 1, e = !0) {
     throw new Error(`'k' passed to topk() must be >= 0 but got ${t}`);
   if (t > n)
     throw new Error(`'k' passed to topk() must be <= the last dimension (${n}) but got ${t}`);
-  const s = { x: o }, i = { k: t, sorted: e }, [r, c] = y.runKernel(j, s, i);
-  return { values: r, indices: c };
+  const s = { x: o }, i = { k: t, sorted: e }, [p, r] = I.runKernel(Y, s, i);
+  return { values: p, indices: r };
 }
-const tt = /* @__PURE__ */ E({ topk_: Z });
+const at = /* @__PURE__ */ x({ topk_: ct });
 /**
  * @license
  * Copyright 2018 Google LLC
@@ -99,13 +101,13 @@ const tt = /* @__PURE__ */ E({ topk_: Z });
  * https://opensource.org/licenses/MIT.
  * =============================================================================
  */
-function et(h) {
-  return new q(h);
+function lt(l) {
+  return new W(l);
 }
-function ot(h) {
-  return new T(h);
+function pt(l) {
+  return new H(l);
 }
-class dt extends R {
+class xt extends V {
   wte;
   // Token embeddings
   wpe;
@@ -119,15 +121,15 @@ class dt extends R {
   log = [];
   // Training log
   constructor(t = {}) {
-    super({ gpt: { ...L, ...t }, layerConfig: {} }), this.wte = new K(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = ot({
+    super({ gpt: { ...F, ...t }, layerConfig: {} }), this.wte = new N(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = pt({
       inputDim: this.config.gpt.blockSize,
       outputDim: this.config.gpt.nEmbed,
       name: "positional_embedding",
-      embeddingsInitializer: P({ mean: 0, stddev: 0.02 })
-    }) : (this.ropeCache = new F(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = et({ rate: this.config.gpt.dropout }), this.blocks = [];
+      embeddingsInitializer: R({ mean: 0, stddev: 0.02 })
+    }) : (this.ropeCache = new A(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = lt({ rate: this.config.gpt.dropout }), this.blocks = [];
     for (let e = 0; e < this.config.gpt.nLayer; e++)
-      this.blocks.push(new v(e, this.config, this));
-    this.lnF = new D(this.config, "final_rms_norm", this);
+      this.blocks.push(new O(e, this.config, this));
+    this.lnF = new G(this.config, "final_rms_norm", this);
   }
   get checkpointing() {
     return this.config.layerConfig.checkpointing === !0;
@@ -136,11 +138,11 @@ class dt extends R {
     this.config.layerConfig.checkpointing = t;
   }
   inputPhase(t, e, o = !1) {
-    return w(() => {
+    return z(() => {
       const n = this.wte.embed(t);
       if (this.config.gpt.useRope === !1) {
-        const [, s] = t.shape, i = this.config.gpt.blockSize, r = H(0, s, 1, "int32"), c = X(A(r, z(e, "int32")), z(i, "int32")), l = this.wpe.apply(c), a = n.add(l);
-        return this.drop.apply(a, { training: o });
+        const [, s] = t.shape, i = this.config.gpt.blockSize, p = et(0, s, 1, "int32"), r = it(Z(p, L(e, "int32")), L(i, "int32")), u = this.wpe.apply(r), c = n.add(u);
+        return this.drop.apply(c, { training: o });
       } else
         return this.drop.apply(n, { training: o });
     });
@@ -167,7 +169,7 @@ class dt extends R {
   }
   calculateLoss(t, e) {
     try {
-      return N()(t, e).mean();
+      return B()(t, e).mean();
     } catch (o) {
       throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
     }
@@ -205,7 +207,7 @@ class dt extends R {
           });
       }*/
   forward(t, e, o) {
-    return this.validateInput(e), w(() => {
+    return this.validateInput(e), z(() => {
       this.startMemory();
       const n = t.cache?.[0]?.length ?? 0;
       let s = this.inputPhase(e, n, t.training);
@@ -213,61 +215,72 @@ class dt extends R {
         throw console.error("Cache", t.cache), new Error(
           `Cache length ${t.cache.length} does not match number of blocks ${this.blocks.length}`
         );
-      for (let c = 0; c < this.blocks.length; c++) {
-        const l = this.blocks[c], a = Math.random() * 1e9, u = {
+      for (let r = 0; r < this.blocks.length; r++) {
+        const u = this.blocks[r], c = Math.random() * 1e9, g = {
           training: t.training,
-          seed: a,
+          seed: c,
           attentionScores: t.attentionScores,
-          pastKV: t.cache ? t.cache[c] : void 0
-        }, p = this.config.layerConfig.checkpointing && t.training ? l.callCheckpoint(u, s) : l.call(u, s);
-        s.dispose(), s = p;
+          pastKV: t.cache ? t.cache[r] : void 0
+        }, S = this.config.layerConfig.checkpointing && t.training ? u.callCheckpoint(g, s) : u.call(g, s);
+        s.dispose(), s = S;
       }
       s = this.lnF.call(t, s);
       const i = this.wte.project(s);
       s.dispose();
-      let r;
-      return o && (r = this.calculateLoss(i, o)), this.endMemory("Forward"), r ? [i, r] : [i];
+      let p;
+      return o && (p = this.calculateLoss(i, o)), this.endMemory("Forward"), p ? [i, p] : [i];
     });
   }
   generate(t, e, o) {
-    const n = o?.temperature ?? 1, s = o?.topK, i = o?.usePadding ?? !1;
-    return w(() => {
-      const r = t, c = r.shape[1], l = c <= this.config.gpt.blockSize ? r : r.slice(
-        [0, c - this.config.gpt.blockSize],
+    const n = o?.temperature ?? 1, s = o?.topK, i = o?.topP, p = o?.usePadding ?? !1;
+    return z(() => {
+      const r = t, u = r.shape[1], c = u <= this.config.gpt.blockSize ? r : r.slice(
+        [0, u - this.config.gpt.blockSize],
         [r.shape[0], this.config.gpt.blockSize]
-      ), a = i ? this.config.gpt.blockSize - l.shape[1] : 0, u = a > 0 ? _(l, [
+      ), g = p ? this.config.gpt.blockSize - c.shape[1] : 0, S = g > 0 ? J(c, [
         [0, 0],
-        [0, a]
-      ]) : l, p = {
+        [0, g]
+      ]) : c, f = {
         training: !1,
         attentionScores: o?.attentionScores ? {
           attentionOut: []
         } : void 0,
         cache: e
-      }, [f] = this.forward(p, u), S = f.shape[1] - 1 - a, M = f.slice([0, S, 0], [f.shape[0], 1, f.shape[2]]);
-      p.attentionScores?.attentionOut && p.attentionScores.attentionOut.forEach((g, k) => {
-        g.shape[1] !== 1 && (p.attentionScores.attentionOut[k] = W(
-          g.slice([0, S, 0], [g.shape[0], 1, g.shape[2]])
-        ), g.dispose());
-      }), f.dispose();
-      const b = M.div(n);
+      }, [d] = this.forward(f, S), M = d.shape[1] - 1 - g, K = d.slice([0, M, 0], [d.shape[0], 1, d.shape[2]]);
+      f.attentionScores?.attentionOut && f.attentionScores.attentionOut.forEach((h, b) => {
+        h.shape[1] !== 1 && (f.attentionScores.attentionOut[b] = tt(
+          h.slice([0, M, 0], [h.shape[0], 1, h.shape[2]])
+        ), h.dispose());
+      }), d.dispose();
+      const w = K.div(n);
       let m;
-      if (s) {
-        const { values: g, indices: k } = tt(b, s), x = I(g.squeeze([1]), 1);
-        m = J(k.squeeze([1]), x, 1);
+      if (i) {
+        const h = q(w.squeeze([1])), b = h.arraySync()[0];
+        h.dispose();
+        const E = b.map((a, k) => ({ prob: a, index: k })).sort((a, k) => k.prob - a.prob);
+        let v = 0;
+        const $ = new Array(E.length).fill(0);
+        for (const a of E)
+          if (v += a.prob, $[a.index] = a.prob, v >= i)
+            break;
+        const _ = $.reduce((a, k) => a + k, 0), D = $.map((a) => a / _);
+        m = C(ot(D), 1, void 0, !0);
+      } else if (s) {
+        const { values: h, indices: b } = at(w, s), E = C(h.squeeze([1]), 1);
+        m = st(b.squeeze([1]), E, 1);
       } else
-        m = I(b.squeeze([1]), 1);
-      let $;
-      return o?.includeProbabilities && ($ = Q(b.squeeze([1]))), m = m.reshape([1, 1]), { output: m, probabilities: $, attention: p.attentionScores?.attentionOut };
+        m = C(w.squeeze([1]), 1);
+      let P;
+      return o?.includeProbabilities && (P = q(w.squeeze([1]))), m = m.reshape([1, 1]), { output: m, probabilities: P, attention: f.attentionScores?.attentionOut };
     });
   }
   getNumParams() {
-    return O(this.config.gpt);
+    return j(this.config.gpt);
   }
   dispose() {
     this.wte.dispose(), this.wpe && this.wpe.dispose(), this.drop.dispose(), this.blocks.forEach((t) => t.dispose()), this.lnF.dispose();
   }
 }
 export {
-  dt as default
+  xt as default
 };

package/dist/{Reshape-BE5rA4rT.js → Reshape-Bt_t7RNz.js}

@@ -1,5 +1,5 @@
-import { ad as $, ae as g, p, af as C, k as x } from "./index-iNhkcAEQ.js";
-import { u as I } from "./gpgpu_math-C0zyxKFi.js";
+import { ad as $, ae as g, q as p, af as C, l as x } from "./index-CnHyhpKc.js";
+import { u as I } from "./gpgpu_math-Df7gzJWH.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -201,12 +201,12 @@ function D(t, e, o) {
  * limitations under the License.
  * =============================================================================
  */
-function U(t) {
+function k(t) {
   const { inputs: e, backend: o, attrs: s } = t, { x: n } = e, { shape: r } = s, i = o, u = p(n.shape), a = C(r, u), c = p(a);
   x(u === c, () => `The new shape (${a}) has ${c} elements and the old shape (${n.shape}) has ${u} elements. The new shape and old shape must have the same number of elements.`);
   const l = i.texData.get(n.dataId);
   return l.isPacked && !f(n.shape, a) && !(l.texture !== null && f(l.shape, a)) ? D(n, a, i) : (i.incRef(n.dataId), { dataId: n.dataId, shape: a, dtype: n.dtype });
 }
 export {
-  U as r
+  k as r
 };

package/dist/TeachableLLM.js

@@ -11,7 +11,7 @@ import g from "./tokeniser/bpe.js";
 import "./papaparse.min-C8l2Kvo1.js";
 import "./index-Tf7vU29b.js";
 import "./jszip.min-CjP2V1VV.js";
-import "./index-iNhkcAEQ.js";
+import "./index-CnHyhpKc.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";

package/dist/TiedEmbedding-DORsPlNL.js

@@ -0,0 +1,44 @@
+import { R as a } from "./random_width-DI2h9CMs.js";
+import "./index-CnHyhpKc.js";
+import { d as s } from "./tfjs_backend-DX9yVvwk.js";
+import o from "./layers/BaseLayer.js";
+import { v as m } from "./variable-BGvK-VN3.js";
+import { g as d } from "./gather-BWyutxwi.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC
+ *
+ * Use of this source code is governed by an MIT-style
+ * license that can be found in the LICENSE file or at
+ * https://opensource.org/licenses/MIT.
+ * =============================================================================
+ */
+function n(e) {
+  return new a(e);
+}
+class c extends o {
+  vocabSize;
+  embedDim;
+  initializer;
+  WEIGHTS;
+  constructor(t, i, r) {
+    super(t, r), this.WEIGHTS = i, this.vocabSize = t.gpt.vocabSize, this.embedDim = t.gpt.nEmbed, this.initializer = n({
+      mean: 0,
+      stddev: 0.02
+    }), this.addVariable(this.WEIGHTS, m(this.initializer.apply([this.vocabSize, this.embedDim]), !0));
+  }
+  embed(t) {
+    return d(this.getVariable(this.WEIGHTS), t, 0);
+  }
+  project(t) {
+    return s(t, this.getVariable(this.WEIGHTS).transpose());
+  }
+  // Dummy, should not be used.
+  forward(t, i) {
+    return this.project(i);
+  }
+}
+export {
+  c as T,
+  n as r
+};

package/dist/{axis_util-97KkkyRQ.js → axis_util-CVbf1vmL.js}

@@ -1,4 +1,4 @@
-import { k as c } from "./index-iNhkcAEQ.js";
+import { l as c } from "./index-CnHyhpKc.js";
 /**
  * @license
  * Copyright 2017 Google LLC. All Rights Reserved.
@@ -21,7 +21,7 @@ function i(e, n) {
       return !1;
   return !0;
 }
-function p(e, n, t) {
+function l(e, n, t) {
   const r = e.length + n.length, s = [];
   let o = 0, f = 0;
   for (let u = 0; u < r; u++)
@@ -37,7 +37,7 @@ function a(e, n) {
 }
 function m(e, n) {
   const t = n.map((r) => 1);
-  return p(e, t, n);
+  return l(e, t, n);
 }
 function d(e, n, t) {
   c(i(n, t), () => `${e} supports only inner-most axes for now. Got axes ${n} and rank-${t} input.`);

package/dist/{broadcast_to-CMlkG8NS.js → broadcast_to-BBoMQXbL.js}

@@ -1,5 +1,5 @@
-import { o as h, i as f, l as p, x as g, E as u, T } from "./index-iNhkcAEQ.js";
-import { r as b } from "./reshape-DxTPgnwL.js";
+import { o as h, j as f, n as p, y as g, E as u, L as b } from "./index-CnHyhpKc.js";
+import { r as T } from "./reshape-CTIbqjwm.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -25,7 +25,7 @@ function m(e, r) {
     const t = n.shape.slice();
     for (; t.length < r.length; )
       t.unshift(1);
-    n = b(n, t);
+    n = T(n, t);
   }
   const s = n.shape, o = Array.from(r);
   for (let t = r.length - 1; t >= 0; t--)
@@ -36,7 +36,7 @@ function m(e, r) {
   if (o.map((t, l) => t > 1 ? l : -1).filter((t) => t >= 0).length === 0)
     return g(n);
   const i = { x: n }, c = { reps: o };
-  return u.runKernel(T, i, c);
+  return u.runKernel(b, i, c);
 }
 const E = /* @__PURE__ */ h({ broadcastTo_: m });
 export {

package/dist/{concat-Cxbo2sOz.js → concat-BRRtq4S2.js}

@@ -1,4 +1,4 @@
-import { o as s, k as a, j as p, x as i, E as l, C as f } from "./index-iNhkcAEQ.js";
+import { o as s, l as a, k as p, y as i, E as l, C as f } from "./index-CnHyhpKc.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.