@genai-fi/nanogpt 0.7.3 → 0.8.1

package/dist/{kernel_funcs_utils-BYKWV8Aa.js → kernel_funcs_utils-CDfFpUab.js}

@@ -1,6 +1,6 @@
-import { k as B, j as G, am as K, a6 as W, an as z, ao as V, ac as N, ap as F, u as S } from "./index-BoWRt-10.js";
-import { u as O, f as Y } from "./gpgpu_math-DGNLNL4I.js";
-import { f as v } from "./backend_util-TE7aTPhZ.js";
+import { l as B, j as G, az as K, aa as z, at as W, aA as V, ag as N, au as F, u as S } from "./index-DdmHGZjq.js";
+import { u as O, f as Y } from "./gpgpu_math-D_ODOLix.js";
+import { f as v } from "./backend_util-yC3YH1jo.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -295,7 +295,7 @@ function L(t) {
   return o.complexTensorInfos = { real: i, imag: a }, n;
 }
 const me = {
-  kernelName: W,
+  kernelName: z,
   backendName: "webgl",
   kernelFunc: L
 };
@@ -315,16 +315,16 @@ const me = {
  * limitations under the License.
  * =============================================================================
  */
-const w = "return (a < 0.) ? b * a : a;", k = `
+const w = "return (a < 0.) ? b * a : a;", R = `
   vec4 aLessThanZero = vec4(lessThan(a, vec4(0.)));
   return (aLessThanZero * (b * a)) + ((vec4(1.0) - aLessThanZero) * a);
 `;
 function oe(t) {
-  const { inputs: e, backend: s, attrs: r } = t, { x: u } = e, { alpha: n } = r, o = s.makeTensorInfo([], "float32", V(n, "float32")), i = N().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new E(k, u.shape, o.shape) : new b(w, u.shape, o.shape), a = s.runWebGLProgram(i, [u, o], "float32");
+  const { inputs: e, backend: s, attrs: r } = t, { x: u } = e, { alpha: n } = r, o = s.makeTensorInfo([], "float32", V(n, "float32")), i = N().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new E(R, u.shape, o.shape) : new b(w, u.shape, o.shape), a = s.runWebGLProgram(i, [u, o], "float32");
   return s.disposeIntermediateTensorInfo(o), a;
 }
 const be = {
-  kernelName: z,
+  kernelName: W,
   backendName: "webgl",
   kernelFunc: oe
 };
@@ -344,12 +344,12 @@ const be = {
  * limitations under the License.
  * =============================================================================
  */
-const R = "return (a < 0.) ? b * a : a;", U = `
+const k = "return (a < 0.) ? b * a : a;", U = `
   vec4 aLessThanZero = vec4(lessThan(a, vec4(0.)));
   return (aLessThanZero * (b * a)) + ((vec4(1.0) - aLessThanZero) * a);
 `;
 function ue(t) {
-  const { inputs: e, backend: s } = t, { x: r, alpha: u } = e, n = N().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new E(U, r.shape, u.shape) : new b(R, r.shape, u.shape);
+  const { inputs: e, backend: s } = t, { x: r, alpha: u } = e, n = N().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new E(U, r.shape, u.shape) : new b(k, r.shape, u.shape);
   return s.runWebGLProgram(n, [r, u], "float32");
 }
 const Ne = {
@@ -386,7 +386,7 @@ function ye({ opSnippet: t, packedOpSnippet: e, cpuKernelImpl: s, dtype: r }) {
     return c ? l = new ne(o.shape, e) : l = new q(o.shape, t), i.runWebGLProgram(l, [o], a);
   };
 }
-function Ie({ opSnippet: t, packedOpSnippet: e, checkOutOfBounds: s = !1, supportsComplex: r = !1, cpuKernelImpl: u, dtype: n }) {
+function Ae({ opSnippet: t, packedOpSnippet: e, checkOutOfBounds: s = !1, supportsComplex: r = !1, cpuKernelImpl: u, dtype: n }) {
   return ({ inputs: o, backend: i }) => {
     const { a, b: c } = o, l = i;
     if (r && a.dtype === "complex64") {
@@ -404,8 +404,8 @@ function Ie({ opSnippet: t, packedOpSnippet: e, checkOutOfBounds: s = !1, suppor
           shape: c.shape
         }, D = new b(t, a.shape, c.shape);
         return l.runWebGLProgram(D, [$, _], S(p.dtype, x.dtype));
-      }), A = L({ inputs: { real: g, imag: m }, backend: l });
-      return l.disposeIntermediateTensorInfo(g), l.disposeIntermediateTensorInfo(m), A;
+      }), I = L({ inputs: { real: g, imag: m }, backend: l });
+      return l.disposeIntermediateTensorInfo(g), l.disposeIntermediateTensorInfo(m), I;
     }
     const d = n || S(a.dtype, c.dtype);
     if ((a.dtype === "string" || c.dtype === "string" || l.shouldExecuteOnCPU([a, c])) && u != null) {
@@ -415,15 +415,15 @@ function Ie({ opSnippet: t, packedOpSnippet: e, checkOutOfBounds: s = !1, suppor
       ) : h, m = a.dtype === "string" ? (
         // tslint:disable-next-line: no-any
         v(f)
-      ) : f, [A, C] = u(a.shape, c.shape, g, m, d), p = l.makeTensorInfo(C, d), x = l.texData.get(p.dataId);
-      return x.values = A, p;
+      ) : f, [I, C] = u(a.shape, c.shape, g, m, d), p = l.makeTensorInfo(C, d), x = l.texData.get(p.dataId);
+      return x.values = I, p;
     }
     const y = N().getBool("WEBGL_PACK_BINARY_OPERATIONS") && e != null;
-    let I;
-    return y ? I = new E(e, a.shape, c.shape, s) : I = new b(t, a.shape, c.shape), l.runWebGLProgram(I, [a, c], d);
+    let A;
+    return y ? A = new E(e, a.shape, c.shape, s) : A = new b(t, a.shape, c.shape), l.runWebGLProgram(A, [a, c], d);
   };
 }
-function Ae(t, e = !1) {
+function Ie(t, e = !1) {
   if (t === "linear")
     return e ? ee : j;
   if (t === "relu")
@@ -433,9 +433,9 @@ function Ae(t, e = !1) {
   if (t === "relu6")
     return e ? ae : Q;
   if (t === "prelu")
-    return e ? U : R;
+    return e ? U : k;
   if (t === "leakyrelu")
-    return e ? k : w;
+    return e ? R : w;
   if (t === "sigmoid")
     return e ? re : X;
   throw new Error(`Activation ${t} has not been implemented for the WebGL backend.`);
@@ -446,7 +446,7 @@ export {
   T as C,
   ne as U,
   Z as a,
-  Ie as b,
+  Ae as b,
   pe as c,
   he as d,
   q as e,
@@ -457,7 +457,7 @@ export {
   fe as j,
   xe as k,
   Oe as l,
-  Ae as m,
+  Ie as m,
   me as n,
   ge as o,
   be as p,

package/dist/layers/BaseLayer.d.ts

@@ -1,27 +1,22 @@
-import { GPTConfig } from '../config';
+import { GPTConfig } from '../models/config';
 import { default as MemoryProfiler } from '../utilities/profile';
 import { default as RoPECache } from './RoPECache';
 import { Tensor, Variable } from '@tensorflow/tfjs-core';
-export interface LayerConfig {
-    checkpointing?: boolean;
-    profiler?: MemoryProfiler;
-    ropeCache?: RoPECache;
-}
-export interface GPTLayerConfig {
-    gpt: GPTConfig;
-    layerConfig: LayerConfig;
-}
 export interface ForwardAttributes {
     training: boolean;
+    checkpointing?: boolean;
+    ropeCache?: RoPECache;
 }
 export default abstract class BaseLayer<ATTR extends ForwardAttributes = ForwardAttributes> {
     readonly parent?: BaseLayer;
-    readonly config: GPTLayerConfig;
+    readonly config: GPTConfig;
     private _variables;
     private _trainable;
     readonly children: BaseLayer[];
-    constructor(config: GPTLayerConfig, parent?: BaseLayer);
+    private profiler?;
+    constructor(config: GPTConfig, parent?: BaseLayer);
     getProfiler(): MemoryProfiler | undefined;
+    setProfiler(profiler: MemoryProfiler | null): void;
     startMemory(): void;
     endMemory(label: string): void;
     addVariable(name: string, variable?: Variable): void;
@@ -29,7 +24,7 @@ export default abstract class BaseLayer<ATTR extends ForwardAttributes = Forward
     get trainableVariables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);
-    getVariable(name: string): Variable;
+    getVariable(name: string, recursive?: boolean): Variable;
     hasVariable(name: string): boolean;
     setVariable(name: string, variable: Variable): void;
     saveWeights(map: Map<string, Tensor[]>): void;

package/dist/layers/BaseLayer.js

@@ -1,22 +1,28 @@
-import { T as g, y as p, e as o, A as v } from "../index-BoWRt-10.js";
-import { v as _ } from "../variable-BuddVFLa.js";
-class M {
+import { T as p, J as g, e as o, K as v } from "../index-DdmHGZjq.js";
+import { v as _ } from "../variable-DPFOJyRG.js";
+class T {
   parent;
   config;
   _variables = /* @__PURE__ */ new Map();
   _trainable = !0;
   children = [];
+  profiler;
   constructor(t, r) {
     this.config = t, this.parent = r, this.parent && this.parent.children.push(this);
   }
   getProfiler() {
-    return this.config.layerConfig.profiler;
+    return this.profiler;
+  }
+  setProfiler(t) {
+    this.profiler = t || void 0, this.children.forEach((r) => {
+      r.setProfiler(t);
+    });
   }
   startMemory() {
-    this.config.layerConfig.profiler?.startMemory();
+    this.profiler?.startMemory();
   }
   endMemory(t) {
-    this.config.layerConfig.profiler?.endMemory(t);
+    this.profiler?.endMemory(t);
   }
   addVariable(t, r) {
     this._variables.set(t, r || null);
@@ -41,11 +47,17 @@ class M {
       r.trainable = t;
     });
   }
-  getVariable(t) {
-    const r = this._variables.get(t);
-    if (!r)
+  getVariable(t, r = !1) {
+    const e = this._variables.get(t);
+    if (!e && r)
+      for (const i of this.children) {
+        const s = i.getVariable(t, !0);
+        if (s)
+          return s;
+      }
+    if (!e)
       throw new Error(`Variable ${t} not found`);
-    return r;
+    return e;
   }
   hasVariable(t) {
     return this._variables.get(t) !== null;
@@ -85,7 +97,7 @@ class M {
   call(t, ...r) {
     this.build();
     const e = this.forward(t, ...r);
-    if (t.training && e instanceof g) {
+    if (t.training && e instanceof p) {
       const i = this.dropout(e);
       return i !== e && e.dispose(), i;
     } else
@@ -95,7 +107,7 @@ class M {
     return this.build(), this.checkpointingFn(t, ...r);
   }
   checkpointingFn(t, ...r) {
-    const e = this.trainableVariables, s = p((...a) => {
+    const e = this.trainableVariables, s = g((...a) => {
       const l = a[a.length - 1], n = a.slice(0, r.length), h = this.forward(t, ...n);
       return l(n), { value: h, gradFunc: (c, f) => {
         const u = o().state.activeTape;
@@ -112,5 +124,5 @@ class M {
   }
 }
 export {
-  M as default
+  T as default
 };

package/dist/layers/CausalSelfAttention.d.ts

@@ -1,5 +1,6 @@
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { Tensor } from '@tensorflow/tfjs-core';
+import { GPTConfig } from '../models/config';
 export type KVCache = {
     k?: Tensor;
     v?: Tensor;
@@ -22,7 +23,7 @@ export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttri
     private projUnits;
     private ATTN;
     private PROJ;
-    constructor(index: number, config: GPTLayerConfig, parent?: BaseLayer);
+    constructor(index: number, config: GPTConfig, parent?: BaseLayer);
     protected build(): void;
     private getAttentionScores;
     private getAttentionScoresWithPast;

package/dist/layers/CausalSelfAttention.js

@@ -3,14 +3,14 @@ import O from "./BaseLayer.js";
 import { qkv as P } from "../ops/qkv.js";
 import { rope as v } from "../ops/rope.js";
 import { appendCache as V } from "../ops/appendCache.js";
-import { w as c, t as C } from "../index-BoWRt-10.js";
+import { k as c, t as C } from "../index-DdmHGZjq.js";
 import { fusedSoftmax as T } from "../ops/fusedSoftmax.js";
-import { d as y } from "../random_width-sZORGo5k.js";
-import { v as b } from "../variable-BuddVFLa.js";
-import { r as k, d as L } from "../dropout-DYs5QFGQ.js";
-import { r as N } from "../reshape-CdBq1WJ6.js";
-import { m as R } from "../mat_mul-8m8pfdcx.js";
-class W extends O {
+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 {
   divisor;
   index;
   units;
@@ -18,27 +18,27 @@ class W extends O {
   ATTN;
   PROJ;
   constructor(t, i, s) {
-    super(i, s), this.index = t, this.units = i.gpt.nEmbed * 3, this.projUnits = i.gpt.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(i.gpt.nEmbed / i.gpt.nHead);
+    super(i, s), this.index = t, this.units = i.nEmbed * 3, this.projUnits = i.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(i.nEmbed / i.nHead);
   }
   build() {
     this.hasVariable(this.ATTN) === !1 && this.setVariable(
       this.ATTN,
       b(
-        k([this.config.gpt.nEmbed, this.units], 0, 0.02),
+        k([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.gpt.nEmbed], 0, 0.02),
+        k([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.gpt.dropout : 0, o);
+    const e = g(t, i, this.divisor), n = T(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.
@@ -47,50 +47,50 @@ class W extends O {
     return o.dispose(), e;
   }
   getQKV(t) {
-    return P(t, this.getVariable(this.ATTN), this.config.gpt.nHead);
+    return P(t, this.getVariable(this.ATTN), this.config.nHead);
   }
   getOutputProjection(t) {
-    const i = t.shape[0], s = t.shape[2], o = this.config.gpt.nEmbed, e = t.transpose([0, 2, 1, 3]), n = N(e, [i, s, o]), p = y(n, this.getVariable(this.PROJ));
-    return n.dispose(), e.dispose(), p;
+    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));
+    return n.dispose(), e.dispose(), r;
   }
   updateCache(t, i, s) {
-    const o = this.config.gpt.blockSize, e = t.shape[2], n = s.length || 0, p = V(t, o, n, s.k);
+    const o = this.config.blockSize, e = t.shape[2], n = s.length || 0, r = V(t, o, n, s.k);
     t.dispose(), s.k && s.k.dispose();
-    const a = V(i, o, n, s.v);
+    const p = V(i, o, n, s.v);
     i.dispose(), s.v && s.v.dispose();
     const d = Math.min(n + e, o), h = s.cumulativeLength + e;
-    s.length = d, s.cumulativeLength = h, s.k = c(p), s.v = c(a);
+    s.length = d, s.cumulativeLength = h, s.k = c(r), s.v = c(p);
   }
   forward(t, i) {
     return C(() => {
       this.startMemory();
-      const [s, o, e] = this.getQKV(i), n = t.pastKV ? t.pastKV.cumulativeLength : 0, p = this.config.layerConfig.ropeCache, a = p ? v(s, p, n) : s, d = p ? v(o, p, n) : o;
-      p && (s.dispose(), o.dispose());
+      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;
-      let r;
-      h > 0 ? r = this.getAttentionScoresWithPast(a, u, h) : r = this.getAttentionScores(a, u, t.training, t.seed || 0), a.dispose(), t.pastKV || u.dispose();
-      const l = R(r, m), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
-      f || r.dispose(), t.pastKV || m.dispose();
+      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 K = r.shape[1], S = r.shape[2];
+        const K = a.shape[1], S = a.shape[2];
         t.attentionScores.attentionOut?.push(
-          c(r.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
+          c(a.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
         );
       }
       return this.endMemory("CausalSelfAttention"), A;
     });
   }
   dropout(t) {
-    if (this.config.gpt.dropout > 0) {
-      const i = L(t, this.config.gpt.dropout);
+    if (this.config.dropout > 0) {
+      const i = y(t, this.config.dropout);
       return t.dispose(), i;
     } else
       return t;
   }
 }
 export {
-  W as default
+  $ as default
 };

package/dist/layers/MLP.d.ts

@@ -1,11 +1,12 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../main';
 export default class MLP extends BaseLayer {
     private index;
     private hiddenUnits;
     private MLPHIDDEN;
     private MLPOUT;
-    constructor(index: number, config: GPTLayerConfig, parent?: BaseLayer);
+    constructor(index: number, config: GPTConfig, parent?: BaseLayer);
     protected build(): void;
     forward(_: ForwardAttributes, x: Tensor): Tensor;
     protected dropout(x: Tensor): Tensor;

package/dist/layers/MLP.js

@@ -1,56 +1,52 @@
-import { t as l } from "../index-BoWRt-10.js";
+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-BuddVFLa.js";
-import { r as h, d as f } from "../dropout-DYs5QFGQ.js";
-import { r as d } from "../reshape-CdBq1WJ6.js";
-import { m as c } from "../mat_mul-8m8pfdcx.js";
-class V extends u {
+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";
+class H extends u {
   index;
   hiddenUnits;
   MLPHIDDEN;
   MLPOUT;
   constructor(i, t, s) {
-    super(t, s), this.index = i, this.hiddenUnits = t.gpt.mlpFactor * t.gpt.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
+    super(t, s), this.index = i, this.hiddenUnits = t.mlpFactor * t.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
   }
   build() {
     this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(
       this.MLPHIDDEN,
       o(
-        h([this.config.gpt.nEmbed, this.hiddenUnits], 0, 0.02),
+        h([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
         !0
         //`block_${this.index}_attn_cAttn_kernel`
       )
     ), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(
       this.MLPOUT,
       o(
-        h(
-          [this.hiddenUnits, this.config.gpt.nEmbed],
-          0,
-          0.02 / Math.sqrt(2 * this.config.gpt.nLayer)
-        ),
+        h([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
         !0
         //`block_${this.index}_attn_cProj_kernel`
       )
     );
   }
   forward(i, t) {
-    return l(() => {
+    return p(() => {
       this.startMemory();
-      const [s, r, e] = t.shape, n = d(t, [s * r, e]), a = M(n, this.getVariable(this.MLPHIDDEN)), p = c(a, this.getVariable(this.MLPOUT));
+      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 m = d(p, [s, r, e]);
-      return this.endMemory("MLP"), m;
+      const l = d(m, [s, r, e]);
+      return this.endMemory("MLP"), l;
     });
   }
   dropout(i) {
-    if (this.config.gpt.dropout > 0) {
-      const t = f(i, this.config.gpt.dropout);
+    if (this.config.dropout > 0) {
+      const t = f(i, this.config.dropout);
       return i.dispose(), t;
     }
     return i;
   }
 }
 export {
-  V as default
+  H as default
 };

package/dist/layers/PositionEmbedding.d.ts

@@ -0,0 +1,9 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer } from './BaseLayer';
+import { GPTConfig, ModelForwardAttributes } from '../main';
+export default class PositionEmbedding extends BaseLayer {
+    private wpe?;
+    private drop;
+    constructor(config: GPTConfig, name?: string, parent?: BaseLayer);
+    forward(attrs: ModelForwardAttributes, x: Tensor): Tensor;
+}

package/dist/layers/PositionEmbedding.js

@@ -0,0 +1,45 @@
+import { t as c, a9 as u, b as i } from "../index-DdmHGZjq.js";
+import f from "./BaseLayer.js";
+import { E as g, D as h } from "../random_width-DKGeiFuR.js";
+import { r as b } from "../exports_initializers-DKk7-bsx.js";
+import { m as l } from "../mod-CbibJi3D.js";
+import { r as w } from "../range-BcUvLuf5.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 E(t) {
+  return new h(t);
+}
+function x(t) {
+  return new g(t);
+}
+class q extends f {
+  wpe;
+  // Position embeddings
+  drop;
+  // Dropout
+  constructor(o, n = "", r) {
+    super(o, r), this.wpe = x({
+      inputDim: this.config.blockSize,
+      outputDim: this.config.nEmbed,
+      name: n,
+      embeddingsInitializer: b({ mean: 0, stddev: 0.02 })
+    }), this.drop = E({ rate: this.config.dropout });
+  }
+  forward(o, n) {
+    const r = o.cache?.[0]?.length ?? 0;
+    return c(() => {
+      const [, s] = n.shape, e = this.config.blockSize, a = w(0, s, 1, "int32"), m = l(u(a, i(r, "int32")), i(e, "int32")), d = this.wpe.apply(m), p = n.add(d);
+      return this.drop.apply(p, { training: o.training });
+    });
+  }
+}
+export {
+  q as default
+};

package/dist/layers/RMSNorm.d.ts

@@ -1,7 +1,8 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../main';
 export default class RMSNorm extends BaseLayer {
     private GAMMA;
-    constructor(config: GPTLayerConfig, name?: string, parent?: BaseLayer);
+    constructor(config: GPTConfig, name?: string, parent?: BaseLayer);
     forward(_: ForwardAttributes, x: Tensor): Tensor;
 }

package/dist/layers/RMSNorm.js

@@ -1,12 +1,12 @@
-import { t as s } from "../index-BoWRt-10.js";
+import { t as s } from "../index-DdmHGZjq.js";
 import e from "./BaseLayer.js";
 import { normRMS as a } from "../ops/normRMS.js";
-import { v as i } from "../variable-BuddVFLa.js";
-import { o as m } from "../ones-Dj0SDhHf.js";
-class f extends e {
+import { v as i } from "../variable-DPFOJyRG.js";
+import { o as m } from "../ones-BAqVh-eA.js";
+class l extends e {
   GAMMA;
   constructor(r, t = "", o) {
-    super(r, o), this.GAMMA = t, this.addVariable(this.GAMMA, i(m([r.gpt.nEmbed]), !0, this.GAMMA, "float32"));
+    super(r, o), this.GAMMA = t, this.addVariable(this.GAMMA, i(m([r.nEmbed]), !0, this.GAMMA, "float32"));
   }
   forward(r, t) {
     return s(() => {
@@ -17,5 +17,5 @@ class f extends e {
   }
 }
 export {
-  f as default
+  l as default
 };

package/dist/layers/RoPECache.d.ts

@@ -1,5 +1,5 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { GPTConfig } from '../config';
+import { GPTConfig } from '../models/config';
 export default class RoPECache {
     private readonly config;
     readonly rotaryDim: number;

package/dist/layers/RoPECache.js

@@ -1,7 +1,7 @@
-import { b as t, x as h, t as n, w as p } from "../index-BoWRt-10.js";
-import { r as c } from "../reciprocal-BvGAyKyu.js";
-import { c as f, s as m } from "../sin-BeA3tsEd.js";
-import { r as a } from "../range-CRuAh-gd.js";
+import { b as t, x as h, t as n, k as p } from "../index-DdmHGZjq.js";
+import { r as c } from "../reciprocal-DhDWSKiD.js";
+import { c as f, s as m } from "../sin-CPxad7Am.js";
+import { r as a } from "../range-BcUvLuf5.js";
 class D {
   constructor(o) {
     this.config = o;

package/dist/layers/TiedEmbedding.d.ts

@@ -1,11 +1,12 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../models/config';
 export default class TiedEmbeddingOutputLayer extends BaseLayer {
     private vocabSize;
     private embedDim;
     private initializer;
     private WEIGHTS;
-    constructor(config: GPTLayerConfig, name: string, parent?: BaseLayer);
+    constructor(config: GPTConfig, name: string, parent?: BaseLayer);
     embed(inputs: Tensor): Tensor;
     project(inputs: Tensor): Tensor;
     forward(_: ForwardAttributes, x: Tensor): Tensor;

package/dist/layers/TiedEmbedding.js

@@ -1,9 +1,31 @@
-import "../random_width-sZORGo5k.js";
-import "../index-BoWRt-10.js";
-import { T as e } from "../TiedEmbedding-BxOerUmB.js";
-import "./BaseLayer.js";
-import "../variable-BuddVFLa.js";
-import "../gather-CMMy2KEG.js";
+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";
+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({
+      mean: 0,
+      stddev: 0.02
+    }), this.addVariable(this.WEIGHTS, m(this.initializer.apply([this.vocabSize, this.embedDim]), !0));
+  }
+  embed(i) {
+    return o(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);
+  }
+}
 export {
-  e as default
+  S as default
 };