@genai-fi/nanogpt 0.4.5 → 0.5.0

package/dist/layers/TransformerBlock.js

@@ -1,53 +1,32 @@
-import h from "./CausalSelfAttention.js";
-import o from "./MLP.js";
-import a from "./RMSNorm.js";
-import p from "./BaseLayer.js";
-import { t as d } from "../index--6vO-cOz.js";
-class W extends p {
+import l from "./CausalSelfAttention.js";
+import r from "./MLP.js";
+import o from "./RMSNorm.js";
+import { B as d } from "../BaseLayer-BhrMN8JO.js";
+import { t as p } from "../index-iNhkcAEQ.js";
+class k extends d {
   ln1;
   attn;
   ln2;
   mlp;
   index;
-  _trainable = !0;
   skipped = !1;
-  constructor(t, s) {
-    super(s), this.index = t, this.ln1 = new a(s, `block_${this.index}_rms1`), this.attn = new h(this.index, s), this.ln2 = new a(s, `block_${this.index}_rms2`), this.mlp = new o(this.index, s);
-  }
-  get variables() {
-    return [
-      ...this.ln1.trainableWeights.map((t) => t),
-      ...this.attn.variables,
-      ...this.ln2.trainableWeights.map((t) => t),
-      ...this.mlp.variables
-    ];
-  }
-  get trainable() {
-    return this._trainable;
-  }
-  set trainable(t) {
-    this._trainable = t, this.ln1.trainable = t, this.ln2.trainable = t, this.attn.trainable = t, this.mlp.trainable = t;
-  }
-  saveWeights(t) {
-    this.attn.saveWeights(t), this.mlp.saveWeights(t), t.set(`block_${this.index}_rms1`, this.ln1.getWeights()), t.set(`block_${this.index}_rms2`, this.ln2.getWeights());
-  }
-  loadWeights(t) {
-    this.attn.loadWeights(t), this.mlp.loadWeights(t), this.ln1.setWeights(t.get(`block_${this.index}_rms1`) || []), this.ln2.setWeights(t.get(`block_${this.index}_rms2`) || []);
+  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);
   }
   getMLPOutput(t, s) {
-    const i = this.ln2.apply(t), e = this.mlp.call(i, s);
-    return t.add(e);
+    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;
   }
-  call(t, s = !1, i = !1, e) {
-    return d(() => {
+  forward(t, s) {
+    return p(() => {
       if (this.skipped)
-        return { output: t };
-      const l = this.ln1.apply(t), n = this.attn.call(l, s, i, e), r = t.add(n.output);
-      return {
-        output: this.getMLPOutput(r, s),
-        attention: n.attention,
-        cache: n.presentKV
-      };
+        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);
     });
   }
   dispose() {
@@ -55,5 +34,5 @@ class W extends p {
   }
 }
 export {
-  W as default
+  k as default
 };

package/dist/{log_sum_exp-CiEy1aUe.js → log_sum_exp-CkumwesB.js}

@@ -1,8 +1,8 @@
-import { o as r, h as p, E as u, a6 as h, a7 as E, $, s as S, a8 as d } from "./index--6vO-cOz.js";
-import { e as K } from "./axis_util-QP0LdI1v.js";
-import { m as T } from "./max-BUShNgfh.js";
-import { r as m } from "./reshape-z51Eu-re.js";
-import { s as _ } from "./sum-DdkDf2MG.js";
+import { o as r, i as p, E as u, a7 as E, a8 as h, a1 as S, s as $, a9 as d } from "./index-iNhkcAEQ.js";
+import { e as K } from "./axis_util-97KkkyRQ.js";
+import { m as T } from "./max-CYaAjEEp.js";
+import { r as m } from "./reshape-DxTPgnwL.js";
+import { s as _ } from "./sum-B_92TaHD.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -21,7 +21,7 @@ import { s as _ } from "./sum-DdkDf2MG.js";
  */
 function b(s) {
   const o = { x: p(s, "x", "exp") };
-  return u.runKernel(h, o);
+  return u.runKernel(E, o);
 }
 const N = /* @__PURE__ */ r({ exp_: b });
 /**
@@ -42,7 +42,7 @@ const N = /* @__PURE__ */ r({ exp_: b });
  */
 function v(s) {
   const o = { x: p(s, "x", "log", "float32") };
-  return u.runKernel(E, o);
+  return u.runKernel(h, o);
 }
 const w = /* @__PURE__ */ r({ log_: v });
 /**
@@ -61,13 +61,13 @@ const w = /* @__PURE__ */ r({ log_: v });
  * limitations under the License.
  * =============================================================================
  */
-function A(s, n = null, o = !1) {
-  const a = p(s, "x", "logSumExp"), t = $(n, a.shape), x = T(
-    a,
+function A(s, a = null, o = !1) {
+  const n = p(s, "x", "logSumExp"), t = S(a, n.shape), x = T(
+    n,
     t,
     !0
     /* keepDims */
-  ), i = S(a, x), l = N(i), f = _(l, t), c = w(f), e = d(m(x, c.shape), c);
+  ), i = $(n, x), l = N(i), f = _(l, t), c = w(f), e = d(m(x, c.shape), c);
   if (o) {
     const g = K(e.shape, t);
     return m(e, g);

package/dist/main.js

@@ -5,7 +5,7 @@ import { default as q } from "./tokeniser/bpe.js";
 import { default as A } from "./utilities/waitForModel.js";
 import { default as I } from "./data/textLoader.js";
 import { estimateMemoryUsage as K, estimateParameterCount as O, estimateResources as Q, estimateTrainingMemoryUsage as S, validateConfig as V } from "./utilities/parameters.js";
-import "./index--6vO-cOz.js";
+import "./index-iNhkcAEQ.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";

package/dist/{mat_mul-BEHRPMh0.js → mat_mul-D0SifYfJ.js}

@@ -1,4 +1,4 @@
-import { o as m, h as s, p as c, E as M, B as p } from "./index--6vO-cOz.js";
+import { o as m, i as s, q as c, E as M, B as p } from "./index-iNhkcAEQ.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -21,7 +21,7 @@ function f(e, o, n = !1, l = !1) {
   const r = { a, b: t }, u = { transposeA: n, transposeB: l };
   return M.runKernel(p, r, u);
 }
-const h = /* @__PURE__ */ m({ matMul_: f });
+const b = /* @__PURE__ */ m({ matMul_: f });
 export {
-  h as m
+  b as m
 };

package/dist/{max-BUShNgfh.js → max-CYaAjEEp.js}

@@ -1,4 +1,4 @@
-import { o as r, h as e, E as x, M as c } from "./index--6vO-cOz.js";
+import { o as r, i as e, E as x, M as c } from "./index-iNhkcAEQ.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -15,11 +15,11 @@ import { o as r, h as e, E as x, M as c } from "./index--6vO-cOz.js";
  * limitations under the License.
  * =============================================================================
  */
-function m(n, o = null, s = !1) {
+function i(n, o = null, s = !1) {
   const t = { x: e(n, "x", "max") }, a = { reductionIndices: o, keepDims: s };
   return x.runKernel(c, t, a);
 }
-const l = /* @__PURE__ */ r({ max_: m });
+const l = /* @__PURE__ */ r({ max_: i });
 export {
   l as m
 };

package/dist/{moments-DYOHXoRV.js → moments-B06NlR_V.js}

@@ -1,6 +1,6 @@
-import { o as m, h as c, E as f, _ as i, $ as l, a0 as h, s as x, x as d } from "./index--6vO-cOz.js";
-import { e as v } from "./axis_util-QP0LdI1v.js";
-import { r as E } from "./reshape-z51Eu-re.js";
+import { o as m, i as c, E as i, a0 as f, a1 as l, a2 as h, s as x, y as d } from "./index-iNhkcAEQ.js";
+import { e as v } from "./axis_util-97KkkyRQ.js";
+import { r as E } from "./reshape-DxTPgnwL.js";
 /**
  * @license
  * Copyright 2020 Google Inc. All Rights Reserved.
@@ -19,7 +19,7 @@ import { r as E } from "./reshape-z51Eu-re.js";
  */
 function S(a, t = null, e = !1) {
   const s = { x: c(a, "x", "mean") }, o = { axis: t, keepDims: e };
-  return f.runKernel(i, s, o);
+  return i.runKernel(f, s, o);
 }
 const r = /* @__PURE__ */ m({ mean_: S });
 /**
@@ -46,8 +46,8 @@ function T(a, t = null, e = !1) {
   const p = h(x(d(a, "float32"), E(s, o))), u = r(p, n, e);
   return { mean: s, variance: u };
 }
-const K = /* @__PURE__ */ m({ moments_: T });
+const N = /* @__PURE__ */ m({ moments_: T });
 export {
-  K as a,
+  N as a,
   r as m
 };

package/dist/{norm-DSva3hI3.js → norm-D3676xIo.js}

@@ -1,8 +1,8 @@
-import { o as l, h as c, E as y, a1 as E, $ as w, a2 as o, a3 as u, U as v, f as I, a0 as $ } from "./index--6vO-cOz.js";
-import { e as A } from "./axis_util-QP0LdI1v.js";
-import { m as f } from "./max-BUShNgfh.js";
-import { r as h } from "./reshape-z51Eu-re.js";
-import { s as t } from "./sum-DdkDf2MG.js";
+import { o as l, i as c, E as y, a3 as E, a1 as w, a4 as o, a5 as u, U as v, f as I, a2 as A } from "./index-iNhkcAEQ.js";
+import { e as $ } from "./axis_util-97KkkyRQ.js";
+import { m as f } from "./max-CYaAjEEp.js";
+import { r as h } from "./reshape-DxTPgnwL.js";
+import { s as t } from "./sum-B_92TaHD.js";
 /**
  * @license
  * Copyright 2020 Google Inc. All Rights Reserved.
@@ -46,7 +46,7 @@ function T(n, e = "euclidean", r = null, m = !1) {
   let i = a.shape;
   if (m) {
     const p = w(r, n.shape);
-    i = A(a.shape, p);
+    i = $(a.shape, p);
   }
   return h(a, i);
 }
@@ -74,7 +74,7 @@ function d(n, e, r = null) {
     if (e === -1 / 0)
       return s(t(o(n), r[1]), r[0]);
     if (e === "fro" || e === "euclidean")
-      return u(t($(n), r));
+      return u(t(A(n), r));
     throw new Error(`Error in norm: invalid ord value: ${e}`);
   }
   throw new Error(`Error in norm: invalid axis: ${r}`);

package/dist/{ones-D6kB8bdY.js → ones-BIeFnPHR.js}

@@ -1,5 +1,5 @@
-import { k as n, l as t, n as m, E as i } from "./index--6vO-cOz.js";
-import { z as l, c } from "./zeros-8xl-W2DC.js";
+import { l as n, n as t, p as m, E as i } from "./index-iNhkcAEQ.js";
+import { z as l, c } from "./zeros-NMYTayy7.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/ops/appendCache.js

@@ -1,12 +1,12 @@
-import { e as a } from "../index--6vO-cOz.js";
+import { e as a } from "../index-iNhkcAEQ.js";
 import "./cpu/appendCache.js";
 import "./webgl/appendCache.js";
-import { z as s } from "../zeros-8xl-W2DC.js";
-import { c } from "../concat-DvWM7HGZ.js";
+import { c as s } from "../concat-Cxbo2sOz.js";
+import { z as c } from "../zeros-NMYTayy7.js";
 function i(r, p, n, o) {
   if (!o) {
     const e = r.shape[2];
-    return c([r, s([r.shape[0], r.shape[1], p - e, r.shape[3]])], 2);
+    return s([r, c([r.shape[0], r.shape[1], p - e, r.shape[3]])], 2);
   }
   return a().runKernel("AppendCache", { cache: o, item: r }, { maxSize: p, pastLen: n });
 }

package/dist/ops/attentionMask.d.ts

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

package/dist/ops/attentionMask.js

@@ -1,10 +1,10 @@
-import { e as i } from "../index--6vO-cOz.js";
+import { e as o } from "../index-iNhkcAEQ.js";
 import "./cpu/attentionMask.js";
 import "./webgl/attentionMask.js";
 import "./grads/attentionMask.js";
-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 });
+function s(t, n, e, r) {
+  return o().runKernel("AttentionMask", { q: t, k: n }, { divisor: e, pastLen: r || 0 });
 }
 export {
-  f as attentionMask
+  s as attentionMask
 };

package/dist/ops/cpu/appendCache.js

@@ -1,5 +1,5 @@
-import { r as d } from "../../index--6vO-cOz.js";
-import { c as h } from "../../concat-DvWM7HGZ.js";
+import { r as d } from "../../index-iNhkcAEQ.js";
+import { c as h } from "../../concat-Cxbo2sOz.js";
 function u(p) {
   const { cache: n, item: s } = p.inputs, { maxSize: r, pastLen: c } = p.attrs, t = n.shape[0], o = n.shape[1], a = n.shape[3], e = s.shape[2];
   if (c + e <= r) {

package/dist/ops/cpu/attentionMask.js

@@ -1,22 +1,21 @@
-import { r as o, f as k } from "../../index--6vO-cOz.js";
-import { m as d } from "../../mat_mul-BEHRPMh0.js";
-function r(t) {
-  const { q: e, k: n, mask: s } = t.inputs, { divisor: c } = t.attrs, m = e.shape[2], i = n.shape[2], a = d(e, n, !1, !0).mul(k(c));
-  if (s) {
-    const l = s.slice([0, 0], [m, i]).expandDims(0).expandDims(0);
-    return a.add(l);
-  }
-  return a;
+import { r as a, g as p, f as u } from "../../index-iNhkcAEQ.js";
+import { l as N, w as b } from "../../ops-ObfXLHYQ.js";
+import { o as g } from "../../ones-BIeFnPHR.js";
+import { z as A } from "../../zeros-NMYTayy7.js";
+import { m as I } from "../../mat_mul-D0SifYfJ.js";
+function o(n) {
+  const { q: s, k: e } = n.inputs, { divisor: r } = n.attrs, c = s.shape[2], t = e.shape[2], m = N.bandPart(g([t, t]), -1, 0).cast("bool"), l = A([t, t]), i = p([t, t], Number.NEGATIVE_INFINITY), f = b(m, l, i), k = I(s, e, !1, !0).mul(u(r)), d = f.slice([0, 0], [c, t]).expandDims(0).expandDims(0);
+  return k.add(d);
 }
-const u = {
+const w = {
   kernelName: "AttentionMask",
   backendName: "cpu",
-  kernelFunc: r
+  kernelFunc: o
 };
-o(u);
-const f = {
+a(w);
+const M = {
   kernelName: "AttentionMask",
   backendName: "tensorflow",
-  kernelFunc: r
+  kernelFunc: o
 };
-o(f);
+a(M);

package/dist/ops/cpu/fusedSoftmax.js

@@ -1,5 +1,5 @@
-import { r as n } from "../../index--6vO-cOz.js";
-import { s as f } from "../../softmax-Dsxflvdl.js";
+import { r as n } from "../../index-iNhkcAEQ.js";
+import { s as f } from "../../softmax-BjsptB07.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, N as h, r as m, s as p } from "../../index--6vO-cOz.js";
-import { r as N } from "../../range-C_vpUjBu.js";
-import { s as l } from "../../stack-CmqSdsfs.js";
+import { o as u, i as c, E as g, L as h, r as m, s as p } from "../../index-iNhkcAEQ.js";
+import { r as l } from "../../range-BsFU-SNG.js";
+import { s as N } from "../../stack--cqr9Dgc.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -23,8 +23,8 @@ function f(e, s) {
 }
 const b = /* @__PURE__ */ u({ gatherND_: f });
 function d(e) {
-  const { values: s, labels: n, logits: r } = e.inputs, t = n.shape[0], o = N(0, t, 1, "int32"), a = l([o, n], 1), i = b(r, a);
-  return p(s, i);
+  const { values: s, labels: n, logits: r } = e.inputs, t = n.shape[0], i = l(0, t, 1, "int32"), o = N([i, n], 1), a = b(r, o);
+  return p(s, a);
 }
 const k = {
   kernelName: "EfficientGatherSub",

package/dist/ops/cpu/gelu.js

@@ -1,4 +1,4 @@
-import { r as t, t as d } from "../../index--6vO-cOz.js";
+import { r as t, t as d } from "../../index-iNhkcAEQ.js";
 const o = 0.7978845608028654, c = 0.044715;
 function m(u) {
   const { inputs: l } = u, { x: e } = l, n = e;

package/dist/ops/cpu/matMulGelu.js

@@ -1,4 +1,4 @@
-import { r as a, t as i } from "../../index--6vO-cOz.js";
+import { r as a, t as i } from "../../index-iNhkcAEQ.js";
 const c = 0.7978845608028654, m = 0.044715;
 function M(o) {
   const { inputs: s } = o, { x: t, kernel: l } = s, e = t, u = l;

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

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

package/dist/ops/cpu/matMulMul.js

@@ -0,0 +1,17 @@
+import { r as n, t as M } from "../../index-iNhkcAEQ.js";
+function e(t) {
+  const { inputs: r, attrs: o } = t, { transposeA: s, transposeB: l } = o, { x: c, kernel: u, y: a } = r, m = c, i = u, k = a;
+  return M(() => m.matMul(i, s, l).mul(k));
+}
+const f = {
+  kernelName: "MatMulMul",
+  backendName: "cpu",
+  kernelFunc: e
+};
+n(f);
+const p = {
+  kernelName: "MatMulMul",
+  backendName: "tensorflow",
+  kernelFunc: e
+};
+n(p);

package/dist/ops/cpu/mulDropout.js

@@ -1,4 +1,4 @@
-import { r as e, b as u } from "../../index--6vO-cOz.js";
+import { r as e, b as u } from "../../index-iNhkcAEQ.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/normRMS.js

@@ -1,4 +1,4 @@
-import { r as o, t as d } from "../../index--6vO-cOz.js";
+import { r as o, t as d } from "../../index-iNhkcAEQ.js";
 function i(t) {
   const { inputs: e } = t, { x: n, gamma: s } = e, r = n, a = s;
   return d(() => {

package/dist/ops/cpu/qkv.js

@@ -1,6 +1,6 @@
-import { r as q } from "../../index--6vO-cOz.js";
-import { r as o } from "../../reshape-z51Eu-re.js";
-import { s as x } from "../../split-B_k_jwud.js";
+import { r as q } from "../../index-iNhkcAEQ.js";
+import { r as o } from "../../reshape-DxTPgnwL.js";
+import { s as x } from "../../split-BCbrzthj.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--6vO-cOz.js";
-import { r as F } from "../../range-C_vpUjBu.js";
-import { g as I } from "../../gather-C5D8PxwA.js";
-import { s as E } from "../../stack-CmqSdsfs.js";
-import { c as T } from "../../concat-DvWM7HGZ.js";
+import { r as S } from "../../index-iNhkcAEQ.js";
+import { r as F } from "../../range-BsFU-SNG.js";
+import { g as I } from "../../gather-Bxe1Qip8.js";
+import { s as E } from "../../stack--cqr9Dgc.js";
+import { c as T } from "../../concat-Cxbo2sOz.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, a5 as w, r as $, s as d, b as m } from "../../index--6vO-cOz.js";
-import { r as b } from "../../range-C_vpUjBu.js";
-import { s as E } from "../../stack-CmqSdsfs.js";
-import { o as D } from "../../ones-D6kB8bdY.js";
+import { o as l, l as g, i, E as k, a6 as w, r as $, s as d, b as m } from "../../index-iNhkcAEQ.js";
+import { r as b } from "../../range-BsFU-SNG.js";
+import { s as E } from "../../stack--cqr9Dgc.js";
+import { o as D } from "../../ones-BIeFnPHR.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)
@@ -51,15 +51,15 @@ function S(n, r, t) {
  * =============================================================================
  */
 function y(n, r, t) {
-  k(t);
-  const s = h(n, "indices", "scatterND", "int32"), e = h(r, "updates", "scatterND");
+  g(t);
+  const s = i(n, "indices", "scatterND", "int32"), e = i(r, "updates", "scatterND");
   S(e, s, t);
   const o = { indices: s, updates: e }, a = { shape: t };
-  return g.runKernel(w, o, a);
+  return k.runKernel(w, o, a);
 }
 const v = /* @__PURE__ */ l({ scatterND_: y });
 function I(n) {
-  const { logits: r, labels: t, dy: s } = n.inputs, e = t.shape[0], o = r.shape[1], a = b(0, e, 1, "int32"), i = E([a, t], 1), c = D([e]), p = v(i, c, [e, o]), f = d(r, p), u = s.reshape([e, 1]);
+  const { logits: r, labels: t, dy: s } = n.inputs, e = t.shape[0], o = r.shape[1], a = b(0, e, 1, "int32"), h = E([a, t], 1), c = D([e]), p = v(h, c, [e, o]), f = d(r, p), u = s.reshape([e, 1]);
   return m(f, u);
 }
 const T = {

package/dist/ops/fusedSoftmax.js

@@ -1,4 +1,4 @@
-import { e as t } from "../index--6vO-cOz.js";
+import { e as t } from "../index-iNhkcAEQ.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--6vO-cOz.js";
+import { e as n } from "../index-iNhkcAEQ.js";
 import "./cpu/gatherSub.js";
 import "./webgl/gatherSub.js";
 function f(r, e, t) {

package/dist/ops/gelu.js

@@ -1,4 +1,4 @@
-import "../index--6vO-cOz.js";
+import "../index-iNhkcAEQ.js";
 import "./cpu/gelu.js";
 import "./webgl/gelu.js";
 import { d as e, g as i } from "./grads/gelu.js";

package/dist/ops/grads/attentionMask.js

@@ -1,21 +1,25 @@
-import { g as i } from "../../index--6vO-cOz.js";
-const u = {
+import { h as l, f as a } from "../../index-iNhkcAEQ.js";
+import { matMulMul as i } from "../matMulMul.js";
+const m = {
   kernelName: "AttentionMask",
   inputsToSave: ["q", "k"],
   outputsToSave: [],
-  gradFunc: (t, a, n) => {
+  gradFunc: (t, u, c) => {
     if (Array.isArray(t))
       throw new Error("Expected dy to be a single Tensor");
-    const [r, e] = a, { divisor: s } = n;
+    const [e, o] = u, { divisor: n } = c;
     return {
-      q: () => t.matMul(e).mul(s),
-      k: () => r.transpose([0, 1, 3, 2]).matMul(t).mul(s).transpose([0, 1, 3, 2]),
+      q: () => i(t, o, a(n)),
+      k: () => {
+        const s = e.transpose([0, 1, 3, 2]), r = i(s, t, a(n));
+        return s.dispose(), r.transpose([0, 1, 3, 2]);
+      },
       mask: () => t,
       divisor: () => {
-        const o = r.matMul(e, !1, !0);
-        return t.mul(o).sum();
+        const s = e.matMul(o, !1, !0), r = t.mul(s);
+        return s.dispose(), r.sum();
       }
     };
   }
 };
-i(u);
+l(m);

package/dist/ops/grads/fusedSoftmax.js

@@ -1,17 +1,20 @@
-import { g as p, b as m, s as d } from "../../index--6vO-cOz.js";
-import { mulDrop as c } from "../mulDrop.js";
-import { s as f } from "../../sum-DdkDf2MG.js";
-const g = {
+import { h as d, b as u, s as f } from "../../index-iNhkcAEQ.js";
+import { mulDrop as l } from "../mulDrop.js";
+import { s as g } from "../../sum-B_92TaHD.js";
+const T = {
   kernelName: "FusedSoftmax",
   outputsToSave: [!0],
-  gradFunc: (s, a, u) => {
-    const [o] = a, { dim: i, dropoutRate: t, seed: r } = u, n = !0, e = t && r ? c(s, o, t, r) : m(s, o);
+  gradFunc: (o, i, n) => {
+    const [s] = i, { dim: a, dropoutRate: t, seed: e } = n, p = !0, r = t && e ? l(o, s, t, e) : u(o, s);
     return {
-      logits: () => d(e, m(f(e, [i], n), o))
+      logits: () => {
+        const m = g(r, [a], p), c = u(m, s);
+        return m.dispose(), f(r, c);
+      }
     };
   }
 };
-p(g);
+d(T);
 export {
-  g as softmaxGradConfig
+  T as softmaxGradConfig
 };

package/dist/ops/grads/gelu.js

@@ -1,4 +1,4 @@
-import { g as t, e as n } from "../../index--6vO-cOz.js";
+import { h as t, e as n } from "../../index-iNhkcAEQ.js";
 import "../cpu/gelu.js";
 import "../webgl/gelu.js";
 const o = {

package/dist/ops/grads/matMulGelu.js

@@ -1,4 +1,4 @@
-import { g as a, e as o } from "../../index--6vO-cOz.js";
+import { h as a, e as o } from "../../index-iNhkcAEQ.js";
 function s(e, n, r) {
   return o().runKernel("MatMulGeluGrad", { dy: e, x: n, kernel: r });
 }

package/dist/ops/grads/normRMS.js

@@ -1,4 +1,4 @@
-import { g as t, e as g } from "../../index--6vO-cOz.js";
+import { h as t, e as g } from "../../index-iNhkcAEQ.js";
 function s(r, a, n) {
   return g().runKernel("RMSNormGrad", { dy: r, x: a, gamma: n });
 }

package/dist/ops/grads/qkv.js

@@ -1,20 +1,30 @@
-import { g as v } from "../../index--6vO-cOz.js";
-const g = {
+import { h as Q } from "../../index-iNhkcAEQ.js";
+const V = {
   kernelName: "QKV",
   inputsToSave: ["x", "kernel"],
   outputsToSave: [],
-  gradFunc: (k, m) => {
-    const [x, K, f] = k, [c, r] = m, [t, s, e] = c.shape, d = x.transpose([0, 2, 1, 3]).reshape([t * s, e]), l = K.transpose([0, 2, 1, 3]).reshape([t * s, e]), u = f.transpose([0, 2, 1, 3]).reshape([t * s, e]), i = r.slice([0, 0], [e, e]), h = r.slice([0, e], [e, e]), M = r.slice([0, 2 * e], [e, e]);
+  gradFunc: (x, K) => {
+    const [f, h, M] = x, [p, l] = K, [t, n, e] = p.shape, i = f.transpose([0, 2, 1, 3]).reshape([t * n, e]), u = h.transpose([0, 2, 1, 3]).reshape([t * n, e]), k = M.transpose([0, 2, 1, 3]).reshape([t * n, e]);
     return {
       x: () => {
-        const n = d.matMul(i, !1, !0), a = l.matMul(h, !1, !0), o = u.matMul(M, !1, !0);
-        return n.add(a).add(o).reshape([t, s, e]);
+        const s = l.slice([0, 0], [e, e]), o = i.matMul(s, !1, !0);
+        s.dispose();
+        const d = l.slice([0, e], [e, e]), r = u.matMul(d, !1, !0);
+        d.dispose();
+        const a = o.add(r);
+        o.dispose(), r.dispose();
+        const c = l.slice([0, 2 * e], [e, e]), m = k.matMul(c, !1, !0);
+        c.dispose();
+        const v = a.add(m).reshape([t, n, e]);
+        return a.dispose(), m.dispose(), v;
       },
       kernel: () => {
-        const n = c.reshape([t * s, e]), a = n.matMul(d, !0, !1), o = n.matMul(l, !0, !1), p = n.matMul(u, !0, !1);
-        return a.concat(o, 1).concat(p, 1);
+        const s = p.reshape([t * n, e]), o = s.matMul(i, !0, !1), d = s.matMul(u, !0, !1), r = o.concat(d, 1);
+        o.dispose(), d.dispose();
+        const a = s.matMul(k, !0, !1), c = r.concat(a, 1);
+        return r.dispose(), a.dispose(), s.dispose(), c;
       }
     };
   }
 };
-v(g);
+Q(V);