@genai-fi/nanogpt 0.4.0 → 0.4.2

package/dist/layers/MLP.js

@@ -1,7 +1,95 @@
-import "../index-XjBAhiFO.js";
-import "./BaseLayer.js";
-import "../random_width-CMHmdbSu.js";
-import { M as i } from "../MLP-KHhikThU.js";
+import { t as _, c as M, e as d, H as v } from "../index-C4JCoBvj.js";
+import x from "./BaseLayer.js";
+import { g as L } from "../gelu-CnCt17Lk.js";
+import { v as n } from "../variable-LJT9Ld63.js";
+import { r as p, d as u } from "../dropout-DfDdklfL.js";
+import { r as l } from "../reshape-Boe4DuIO.js";
+import { m as f } from "../mat_mul-415y5Qn2.js";
+class V extends x {
+  cFc = null;
+  cProj = null;
+  index;
+  _trainable = !0;
+  hiddenUnits;
+  constructor(t, s) {
+    super(s), this.index = t, this.hiddenUnits = s.gpt.mlpFactor * s.gpt.nEmbed;
+  }
+  build() {
+    this.cFc === null && (this.cFc = n(
+      p([this.config.gpt.nEmbed, this.hiddenUnits], 0, 0.02),
+      !0
+      //`block_${this.index}_attn_cAttn_kernel`
+    )), this.cProj === null && (this.cProj = n(
+      p(
+        [this.hiddenUnits, this.config.gpt.nEmbed],
+        0,
+        0.02 / Math.sqrt(2 * this.config.gpt.nLayer)
+      ),
+      !0
+      //`block_${this.index}_attn_cProj_kernel`
+    ));
+  }
+  get variables() {
+    return [this.cFc, this.cProj];
+  }
+  get trainable() {
+    return this._trainable;
+  }
+  set trainable(t) {
+    this._trainable = t, this.cFc && (this.cFc.trainable = t), this.cProj && (this.cProj.trainable = t);
+  }
+  saveWeights(t) {
+    t.set(`block_${this.index}_mlpHidden`, this.cFc ? [this.cFc.clone()] : []), t.set(`block_${this.index}_mlpOut`, this.cProj ? [this.cProj.clone()] : []);
+  }
+  loadWeights(t) {
+    const s = t.get(`block_${this.index}_mlpOut`)?.[0], i = t.get(`block_${this.index}_mlpHidden`)?.[0];
+    if (!s || !i)
+      throw new Error(`Weights for block ${this.index} not found`);
+    this.cFc ? this.cFc.assign(i) : this.cFc = n(i, !0), this.cProj ? this.cProj.assign(s) : this.cProj = n(s, !0);
+  }
+  forward(t) {
+    return _(() => {
+      this.startMemory();
+      const [s, i, o] = t.shape, r = l(t, [s * i, o]), e = f(r, this.cFc), c = L(e);
+      e.dispose();
+      const a = f(c, this.cProj);
+      c.dispose();
+      const h = l(a, [s, i, o]);
+      return this.endMemory("MLP"), h;
+    });
+  }
+  call(t, s = !1) {
+    if (this.build(), s && this.config.layerConfig.checkpointMLP) {
+      const o = M(
+        // @ts-expect-error Invalid params
+        (r, e, c, a) => {
+          const h = this.forward(r);
+          return a([r]), { value: h, gradFunc: (g, m) => {
+            const [b] = m, P = d().state.activeTape;
+            d().state.activeTape = [];
+            const j = v((F, w, T) => this.forward(F))([b, e, c], g);
+            return d().state.activeTape = P, j;
+          } };
+        }
+      )(t, this.cFc, this.cProj);
+      if (this.config.gpt.dropout > 0) {
+        const r = u(o, this.config.gpt.dropout);
+        return o.dispose(), r;
+      }
+      return o;
+    } else {
+      const i = this.forward(t);
+      if (s && this.config.gpt.dropout > 0) {
+        const o = u(i, this.config.gpt.dropout);
+        return i.dispose(), o;
+      }
+      return i;
+    }
+  }
+  dispose() {
+    this.cFc?.dispose(), this.cProj?.dispose();
+  }
+}
 export {
-  i as default
+  V as default
 };

package/dist/layers/RMSNorm.js

@@ -1,7 +1,7 @@
-import { t as r } from "../index-XjBAhiFO.js";
+import { t as r } from "../index-C4JCoBvj.js";
 import m from "./BaseLayer.js";
-import { v as i } from "../variable-BS4AKqNU.js";
-import { o } from "../ones-BO01zpJG.js";
+import { v as i } from "../variable-LJT9Ld63.js";
+import { o } from "../ones-Bf3YR48P.js";
 class d extends m {
   gamma;
   epsilon;

package/dist/layers/RoPECache.js

@@ -1,6 +1,6 @@
-import { o as h, h as c, E as f, N as l, f as n, O as m, t as u, F as p } from "../index-XjBAhiFO.js";
-import { c as d, s as C } from "../sin-BYM-U4Ut.js";
-import { r as a } from "../range-DQMNzBWs.js";
+import { o as h, h as c, E as f, N as l, f as n, O as m, t as u, F as p } from "../index-C4JCoBvj.js";
+import { c as d, s as C } from "../sin-KmhiDuMa.js";
+import { r as a } from "../range-9AzeApCc.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/layers/TiedEmbedding.js

@@ -1,48 +1,8 @@
-import { r as t } from "../random_width-CMHmdbSu.js";
-import "../index-XjBAhiFO.js";
-import { d as s } from "../tfjs_backend-Bzl2SrRo.js";
-import { v as r } from "../variable-BS4AKqNU.js";
-import { g as d } from "../gather-BUmJIS8n.js";
-class n {
-  vocabSize;
-  embedDim;
-  tiedWeights;
-  initializer;
-  constructor(i, e) {
-    this.vocabSize = i.vocabSize, this.embedDim = i.embedDim, this.initializer = t({
-      mean: 0,
-      stddev: 0.02
-    }), this.tiedWeights = r(
-      this.initializer.apply([this.vocabSize, this.embedDim]),
-      !0,
-      e || "tied_embedding"
-    );
-  }
-  get variables() {
-    return [this.tiedWeights];
-  }
-  embed(i) {
-    return d(this.tiedWeights, i, 0);
-  }
-  project(i) {
-    return s(i, this.tiedWeights.transpose());
-  }
-  getWeights() {
-    return [this.tiedWeights];
-  }
-  setWeights(i) {
-    this.tiedWeights.assign(i[0]);
-  }
-  getConfig() {
-    return {
-      vocabSize: this.vocabSize,
-      embedDim: this.embedDim
-    };
-  }
-  dispose() {
-    this.tiedWeights.dispose();
-  }
-}
+import { T as a } from "../TiedEmbedding-CnJ1bx4q.js";
+import "../index-C4JCoBvj.js";
+import "../tfjs_backend-Cug-PH75.js";
+import "../variable-LJT9Ld63.js";
+import "../gather-ZYRWhmXR.js";
 export {
-  n as default
+  a as default
 };

package/dist/layers/TransformerBlock.js

@@ -1,8 +1,8 @@
 import h from "./CausalSelfAttention.js";
-import { M as o } from "../MLP-KHhikThU.js";
+import o from "./MLP.js";
 import a from "./RMSNorm.js";
 import p from "./BaseLayer.js";
-import { t as d } from "../index-XjBAhiFO.js";
+import { t as d } from "../index-C4JCoBvj.js";
 class W extends p {
   ln1;
   attn;

package/dist/{log_sum_exp-DJPkVZZn.js → log_sum_exp-BswFnwOb.js}

@@ -1,8 +1,8 @@
-import { o as r, h as p, E as u, a3 as h, a4 as E, Y as S, s as $, a5 as d } from "./index-XjBAhiFO.js";
-import { e as K } from "./axis_util-DeydwOoC.js";
-import { m as T } from "./max-DJvEiCAJ.js";
-import { r as m } from "./reshape-DFzh97Sc.js";
-import { s as _ } from "./sum-C6u3xMi3.js";
+import { o as r, h as p, E as u, a3 as h, a4 as E, Y as S, s as $, a5 as d } from "./index-C4JCoBvj.js";
+import { e as K } from "./axis_util-BgTGy5w8.js";
+import { m as T } from "./max-CP_9O2Yd.js";
+import { r as m } from "./reshape-Boe4DuIO.js";
+import { s as _ } from "./sum-R28pucR5.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/main.js

@@ -5,7 +5,7 @@ import { default as w } from "./tokeniser/bpe.js";
 import { default as D } from "./utilities/waitForModel.js";
 import { default as F } from "./data/textLoader.js";
 import { estimateMemoryUsage as N, estimateParameterCount as R, estimateResources as j, estimateTrainingMemoryUsage as q, validateConfig as z } from "./utilities/parameters.js";
-import "./index-XjBAhiFO.js";
+import "./index-C4JCoBvj.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";

package/dist/{mat_mul-CKwFEV1Q.js → mat_mul-415y5Qn2.js}

@@ -1,4 +1,4 @@
-import { o as m, h as s, p as c, E as M, B as p } from "./index-XjBAhiFO.js";
+import { o as m, h as s, p as c, E as M, B as p } from "./index-C4JCoBvj.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.

package/dist/{max-DJvEiCAJ.js → max-CP_9O2Yd.js}

@@ -1,4 +1,4 @@
-import { o as r, h as e, E as x, M as c } from "./index-XjBAhiFO.js";
+import { o as r, h as e, E as x, M as c } from "./index-C4JCoBvj.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.

package/dist/{moments-CrWRPcR3.js → moments-CjeIaVdp.js}

@@ -1,6 +1,6 @@
-import { o as m, h as c, E as f, X as i, Y as l, Z as h, s as x, x as d } from "./index-XjBAhiFO.js";
-import { e as v } from "./axis_util-DeydwOoC.js";
-import { r as E } from "./reshape-DFzh97Sc.js";
+import { o as m, h as c, E as f, X as i, Y as l, Z as h, s as x, x as d } from "./index-C4JCoBvj.js";
+import { e as v } from "./axis_util-BgTGy5w8.js";
+import { r as E } from "./reshape-Boe4DuIO.js";
 /**
  * @license
  * Copyright 2020 Google Inc. All Rights Reserved.

package/dist/{norm-BzY929B_.js → norm-CZM380I3.js}

@@ -1,8 +1,8 @@
-import { o as l, h as c, E as y, _ as E, Y as w, $ as o, a0 as u, O as v, f as I, Z as $ } from "./index-XjBAhiFO.js";
-import { e as A } from "./axis_util-DeydwOoC.js";
-import { m as f } from "./max-DJvEiCAJ.js";
-import { r as h } from "./reshape-DFzh97Sc.js";
-import { s as t } from "./sum-C6u3xMi3.js";
+import { o as l, h as c, E as y, _ as E, Y as w, $ as o, a0 as u, O as v, f as I, Z as $ } from "./index-C4JCoBvj.js";
+import { e as A } from "./axis_util-BgTGy5w8.js";
+import { m as f } from "./max-CP_9O2Yd.js";
+import { r as h } from "./reshape-Boe4DuIO.js";
+import { s as t } from "./sum-R28pucR5.js";
 /**
  * @license
  * Copyright 2020 Google Inc. All Rights Reserved.

package/dist/{ones-BO01zpJG.js → ones-Bf3YR48P.js}

@@ -1,5 +1,5 @@
-import { k as n, l as t, n as m, E as i } from "./index-XjBAhiFO.js";
-import { z as l, c } from "./zeros-CmJFiC84.js";
+import { k as n, l as t, n as m, E as i } from "./index-C4JCoBvj.js";
+import { z as l, c } from "./zeros-dnQxFgAD.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/ops/appendCache.js

@@ -1,4 +1,4 @@
-import { e as p } from "../index-XjBAhiFO.js";
+import { e as p } from "../index-C4JCoBvj.js";
 import "./cpu/appendCache.js";
 import "./webgl/appendCache.js";
 function a(e, n, r) {

package/dist/ops/attentionMask.d.ts

@@ -1,2 +1,2 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-export declare function attentionMask(q: Tensor, k: Tensor, mask: Tensor, divisor: number): Tensor;
+export declare function attentionMask(q: Tensor, k: Tensor, divisor: number, mask?: Tensor, pastLen?: number): Tensor;

package/dist/ops/attentionMask.js

@@ -1,10 +1,10 @@
-import { e as o } from "../index-XjBAhiFO.js";
+import { e as i } from "../index-C4JCoBvj.js";
 import "./cpu/attentionMask.js";
 import "./webgl/attentionMask.js";
 import "./grads/attentionMask.js";
-function s(t, n, r, e) {
-  return o().runKernel("AttentionMask", { q: t, k: n, mask: r }, { divisor: e });
+function f(t, n, e, r, o) {
+  return r ? i().runKernel("AttentionMask", { q: t, k: n, mask: r }, { divisor: e, pastLen: o || 0 }) : i().runKernel("AttentionMask", { q: t, k: n }, { divisor: e, pastLen: o || 0 });
 }
 export {
-  s as attentionMask
+  f as attentionMask
 };

package/dist/ops/cpu/appendCache.js

@@ -1,5 +1,5 @@
-import { r as a, e as m } from "../../index-XjBAhiFO.js";
-import { c as d } from "../../concat-DS_qH7MI.js";
+import { r as a, e as m } from "../../index-C4JCoBvj.js";
+import { c as d } from "../../concat-CuRsVY-K.js";
 function r(n) {
   const { cache: c, item: t } = n.inputs, { maxSize: o } = n.attrs, e = d([c, t], 2), s = e.shape[2];
   if (s > o) {

package/dist/ops/cpu/attentionMask.js

@@ -1,18 +1,22 @@
-import { r as s, f as i } from "../../index-XjBAhiFO.js";
-import { m as k } from "../../mat_mul-CKwFEV1Q.js";
-function a(t) {
-  const { q: e, k: o, mask: r } = t.inputs, { divisor: c } = t.attrs, n = e.shape[2], m = k(e, o, !1, !0).mul(i(c)), l = r.slice([0, 0], [n, n]).expandDims(0).expandDims(0);
-  return m.add(l);
+import { r as o, f as l } from "../../index-C4JCoBvj.js";
+import { m as k } from "../../mat_mul-415y5Qn2.js";
+function r(t) {
+  const { q: e, k: c, mask: n } = t.inputs, { divisor: m } = t.attrs, s = e.shape[2], a = k(e, c, !1, !0).mul(l(m));
+  if (n) {
+    const i = n.slice([0, 0], [s, s]).expandDims(0).expandDims(0);
+    return a.add(i);
+  }
+  return a;
 }
 const d = {
   kernelName: "AttentionMask",
   backendName: "cpu",
-  kernelFunc: a
+  kernelFunc: r
 };
-s(d);
+o(d);
 const u = {
   kernelName: "AttentionMask",
   backendName: "tensorflow",
-  kernelFunc: a
+  kernelFunc: r
 };
-s(u);
+o(u);

package/dist/ops/cpu/fusedSoftmax.js

@@ -1,5 +1,5 @@
-import { r as n } from "../../index-XjBAhiFO.js";
-import { s as f } from "../../softmax-4DOn6cPq.js";
+import { r as n } from "../../index-C4JCoBvj.js";
+import { s as f } from "../../softmax-Cujsg4ay.js";
 function r(t) {
   const { inputs: s, attrs: i } = t, { logits: o } = s, { dim: a, dropoutRate: e } = i;
   if (!o)

package/dist/ops/cpu/gatherSub.js

@@ -1,6 +1,6 @@
-import { o as u, h as c, E as g, L as h, r as m, s as p } from "../../index-XjBAhiFO.js";
-import { r as l } from "../../range-DQMNzBWs.js";
-import { s as N } from "../../stack-DaIMO5iX.js";
+import { o as u, h as c, E as g, L as h, r as m, s as p } from "../../index-C4JCoBvj.js";
+import { r as l } from "../../range-9AzeApCc.js";
+import { s as N } from "../../stack-D1YjmgKN.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/ops/cpu/gelu.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/gelu.js

@@ -0,0 +1,40 @@
+import { r as t, t as d } from "../../index-C4JCoBvj.js";
+const o = 0.7978845608028654, c = 0.044715;
+function m(u) {
+  const { inputs: l } = u, { x: e } = l, n = e;
+  return d(() => {
+    const r = n.mul(n).mul(n), s = n.add(r.mul(c)).mul(o).tanh().add(1).mul(0.5);
+    return n.mul(s);
+  });
+}
+const K = {
+  kernelName: "Gelu",
+  backendName: "cpu",
+  kernelFunc: m
+};
+t(K);
+const p = {
+  kernelName: "Gelu",
+  backendName: "tensorflow",
+  kernelFunc: m
+};
+t(p);
+function i(u) {
+  const { dy: l, x: e } = u.inputs;
+  return d(() => {
+    const n = e.square(), r = n.mul(e), a = e.add(r.mul(c)).mul(o).tanh(), f = a.square().neg().add(1), k = n.mul(3 * c).add(1), g = a.add(1).mul(0.5), G = e.mul(f).mul(o).mul(k).mul(0.5), N = g.add(G);
+    return l.mul(N);
+  });
+}
+const x = {
+  kernelName: "GeluGrad",
+  backendName: "cpu",
+  kernelFunc: i
+};
+t(x);
+const h = {
+  kernelName: "GeluGrad",
+  backendName: "tensorflow",
+  kernelFunc: i
+};
+t(h);

package/dist/ops/cpu/mulDropout.js

@@ -1,4 +1,4 @@
-import { r as e, b as u } from "../../index-XjBAhiFO.js";
+import { r as e, b as u } from "../../index-C4JCoBvj.js";
 function n(o) {
   const { inputs: r } = o, { a: l, b: t } = r;
   return console.warn("Using fallback mulDrop implementation without dropout."), u(l, t);

package/dist/ops/cpu/qkv.js

@@ -1,6 +1,6 @@
-import { r as q } from "../../index-XjBAhiFO.js";
-import { r as o } from "../../reshape-DFzh97Sc.js";
-import { s as x } from "../../split-CkbeVdF8.js";
+import { r as q } from "../../index-C4JCoBvj.js";
+import { r as o } from "../../reshape-Boe4DuIO.js";
+import { s as x } from "../../split-DbcNm1-i.js";
 function v(p) {
   const { x: c, kernel: K } = p.inputs, { heads: n } = p.attrs, [s, e, t] = c.shape, a = o(c, [s * e, t]), i = a.dot(K);
   a.dispose();

package/dist/ops/cpu/rope.js

@@ -1,8 +1,8 @@
-import { r as S } from "../../index-XjBAhiFO.js";
-import { r as F } from "../../range-DQMNzBWs.js";
-import { g as I } from "../../gather-BUmJIS8n.js";
-import { s as E } from "../../stack-DaIMO5iX.js";
-import { c as T } from "../../concat-DS_qH7MI.js";
+import { r as S } from "../../index-C4JCoBvj.js";
+import { r as F } from "../../range-9AzeApCc.js";
+import { g as I } from "../../gather-ZYRWhmXR.js";
+import { s as E } from "../../stack-D1YjmgKN.js";
+import { c as T } from "../../concat-CuRsVY-K.js";
 function U(t, c, p, o, r) {
   const n = o.shape[3], s = p;
   if (s > n) return o;

package/dist/ops/cpu/scatterSub.js

@@ -1,7 +1,7 @@
-import { o as l, k, h, E as g, a2 as w, r as $, s as d, b as m } from "../../index-XjBAhiFO.js";
-import { r as b } from "../../range-DQMNzBWs.js";
-import { s as E } from "../../stack-DaIMO5iX.js";
-import { o as D } from "../../ones-BO01zpJG.js";
+import { o as l, k, h, E as g, a2 as w, r as $, s as d, b as m } from "../../index-C4JCoBvj.js";
+import { r as b } from "../../range-9AzeApCc.js";
+import { s as E } from "../../stack-D1YjmgKN.js";
+import { o as D } from "../../ones-Bf3YR48P.js";
 function N(n, r, t) {
   const s = r.rank > 1 ? r.shape[r.rank - 1] : 1, e = r.rank > 1 ? r.rank - 1 : 1, o = `Must have updates.shape = indices.shape[:batchDim] + shape[sliceDim:], got updates.shape: ${t.shape}, indices.shape: ${r.shape}, shape: ${n}, sliceDim: ${s}, and batchDim: ${e}.`;
   if (t.rank < e)

package/dist/ops/fusedSoftmax.js

@@ -1,4 +1,4 @@
-import { e as t } from "../index-XjBAhiFO.js";
+import { e as t } from "../index-C4JCoBvj.js";
 import "./cpu/fusedSoftmax.js";
 import "./webgl/fusedSoftmax.js";
 import "./grads/fusedSoftmax.js";

package/dist/ops/gatherSub.js

@@ -1,4 +1,4 @@
-import { e as n } from "../index-XjBAhiFO.js";
+import { e as n } from "../index-C4JCoBvj.js";
 import "./cpu/gatherSub.js";
 import "./webgl/gatherSub.js";
 function f(r, e, t) {

package/dist/ops/gelu.d.ts

@@ -0,0 +1,3 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function gelu(x: Tensor): Tensor;
+export declare function dGelu(dy: Tensor, x: Tensor): Tensor;

package/dist/ops/gelu.js

@@ -0,0 +1,8 @@
+import "../index-C4JCoBvj.js";
+import "./cpu/gelu.js";
+import "./webgl/gelu.js";
+import { d as e, g as i } from "../gelu-CnCt17Lk.js";
+export {
+  e as dGelu,
+  i as gelu
+};

package/dist/ops/grads/attentionMask.js

@@ -1,4 +1,4 @@
-import { g as i } from "../../index-XjBAhiFO.js";
+import { g as i } from "../../index-C4JCoBvj.js";
 const u = {
   kernelName: "AttentionMask",
   inputsToSave: ["q", "k"],

package/dist/ops/grads/fusedSoftmax.js

@@ -1,6 +1,6 @@
-import { g as p, b as m, s as d } from "../../index-XjBAhiFO.js";
+import { g as p, b as m, s as d } from "../../index-C4JCoBvj.js";
 import { mulDrop as c } from "../mulDrop.js";
-import { s as f } from "../../sum-C6u3xMi3.js";
+import { s as f } from "../../sum-R28pucR5.js";
 const g = {
   kernelName: "FusedSoftmax",
   outputsToSave: [!0],

package/dist/ops/grads/gelu.d.ts

@@ -0,0 +1,2 @@
+import { GradConfig } from '@tensorflow/tfjs-core';
+export declare const geluGradConfig: GradConfig;

package/dist/ops/grads/gelu.js

@@ -0,0 +1,5 @@
+import "../../index-C4JCoBvj.js";
+import { a as e } from "../../gelu-CnCt17Lk.js";
+export {
+  e as geluGradConfig
+};

package/dist/ops/grads/qkv.js

@@ -1,4 +1,4 @@
-import { g as v } from "../../index-XjBAhiFO.js";
+import { g as v } from "../../index-C4JCoBvj.js";
 const g = {
   kernelName: "QKV",
   inputsToSave: ["x", "kernel"],

package/dist/ops/grads/rope.js

@@ -1,4 +1,4 @@
-import { g as a, e as i } from "../../index-XjBAhiFO.js";
+import { g as a, e as i } from "../../index-C4JCoBvj.js";
 function p(n, e, s, o) {
   return i().runKernel("Rope", { x: n, sin: e, cos: s }, { pastLen: o });
 }

package/dist/ops/mulDrop.js

@@ -1,4 +1,4 @@
-import { e as t } from "../index-XjBAhiFO.js";
+import { e as t } from "../index-C4JCoBvj.js";
 import "./cpu/mulDropout.js";
 import "./webgl/mulDropout.js";
 function m(r, o, e, n) {

package/dist/ops/node/sparseCrossEntropy.js

@@ -1,4 +1,4 @@
-import { r as o } from "../../index-XjBAhiFO.js";
+import { r as o } from "../../index-C4JCoBvj.js";
 function r(e) {
   const { logits: t, labels: n } = e.inputs;
   return e.backend.executeMultipleOutputs("SparseSoftmaxCrossEntropyWithLogits", [], [t, n], 2);

package/dist/ops/qkv.js

@@ -1,4 +1,4 @@
-import { e as o } from "../index-XjBAhiFO.js";
+import { e as o } from "../index-C4JCoBvj.js";
 import "./cpu/qkv.js";
 import "./webgl/qkv.js";
 import "./grads/qkv.js";

package/dist/ops/scatterSub.js

@@ -1,4 +1,4 @@
-import { e as i } from "../index-XjBAhiFO.js";
+import { e as i } from "../index-C4JCoBvj.js";
 import "./cpu/scatterSub.js";
 import "./webgl/scatterSub.js";
 function c(t, r, e) {

package/dist/ops/webgl/appendCache.js

@@ -1,4 +1,4 @@
-import { r as h } from "../../index-XjBAhiFO.js";
+import { r as h } from "../../index-C4JCoBvj.js";
 class m {
   variableNames = ["cache", "item"];
   outputShape;

package/dist/ops/webgl/attentionMask.js

@@ -1,11 +1,14 @@
-import { r as m } from "../../index-XjBAhiFO.js";
-class k {
-  variableNames = ["q", "k", "mask"];
+import { r as h } from "../../index-C4JCoBvj.js";
+class l {
+  variableNames = ["q", "k"];
   outputShape;
   userCode;
-  customUniforms = [{ name: "divisor", type: "float" }];
-  constructor(t, a, o, e) {
-    this.outputShape = [t, a, o, o], this.userCode = `
+  customUniforms = [
+    { name: "divisor", type: "float" },
+    { name: "pastLen", type: "int" }
+  ];
+  constructor(t, s, e, n, a) {
+    this.outputShape = [t, s, e, n], this.userCode = `
         void main() {
             ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
             int b = coords.x;
@@ -14,30 +17,28 @@ class k {
             int t2 = coords.w;
 
             float sum = 0.0;
-            for (int i = 0; i < ${e}; ++i) {
+            for (int i = 0; i < ${a}; ++i) {
                 float qv = getQ(b, h, t1, i);
-                float kv = getK(b, h, t2, i); // k is transposed on last two dims
+                float kv = getK(b, h, t2, i);
                 sum += qv * kv;
             }
 
             // Scale by divisor
             float scaled = sum * divisor;
 
-            // Add mask
-            float maskVal = getMask(t1, t2); // mask is [T,T]
-
-            setOutput(scaled + maskVal);
+            // Mask out future positions
+            setOutput((t2 > t1 + pastLen) ? -1.0/0.0 : scaled);
         }
         `;
   }
 }
-function l(s) {
-  const { q: t, k: a, mask: o } = s.inputs, { divisor: e } = s.attrs, n = s.backend, i = t.shape[0], r = t.shape[2], c = t.shape[1], d = new k(i, c, r, t.shape[3]);
-  return n.runWebGLProgram(d, [t, a, o], "float32", [[e]]);
+function m(o) {
+  const { q: t, k: s } = o.inputs, { divisor: e, pastLen: n } = o.attrs, a = o.backend, i = t.shape[0], r = t.shape[2], c = s.shape[2], u = t.shape[1], p = t.shape[3], d = new l(i, u, r, c, p);
+  return a.runWebGLProgram(d, [t, s], "float32", [[e], [n]]);
 }
-const u = {
+const k = {
   kernelName: "AttentionMask",
   backendName: "webgl",
-  kernelFunc: l
+  kernelFunc: m
 };
-m(u);
+h(k);