@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/models/NanoGPTV1.js

@@ -1,12 +1,14 @@
-import { defaultConfig as u } from "./config.js";
-import a from "../layers/TransformerBlock.js";
-import f from "../layers/TiedEmbedding.js";
-import g from "../layers/RoPECache.js";
-import b from "../layers/RMSNorm.js";
-import { t as c, k as p } from "../index-BzFyqcy-.js";
-import w from "./model.js";
-import k from "../layers/PositionEmbedding.js";
-class R extends w {
+import { defaultConfig as g } from "./config.js";
+import b from "../layers/TransformerBlock.js";
+import k from "../layers/TiedEmbedding.js";
+import w from "../layers/RoPECache.js";
+import E from "../layers/RMSNorm.js";
+import { t as l, k as u } from "../index-ZyQhjEPo.js";
+import C from "./model.js";
+import P from "../layers/PositionEmbedding.js";
+import { packingSupported as _ } from "../utilities/packed.js";
+import { p as y, u as M } from "../pack16-CFUqumar.js";
+class I extends C {
   wte;
   // Token embeddings
   wpe;
@@ -17,56 +19,63 @@ class R extends w {
   // Final layer norm
   ropeCache;
   constructor(e = {}) {
-    super({ ...u, ...e }), this.wte = new f(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new k(this.config, "positional_embedding", this) : this.ropeCache = new g(this.config), this.blocks = [];
+    super({ ...g, ...e }), this.wte = new k(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new P(this.config, "positional_embedding", this) : this.ropeCache = new w(this.config), this.blocks = [];
     for (let i = 0; i < this.config.nLayer; i++)
-      this.blocks.push(new a(i, this.config, this));
-    this.lnF = new b(this.config, "final_rms_norm", this);
+      this.blocks.push(new b(i, this.config, this));
+    this.lnF = new E(this.config, "final_rms_norm", this);
   }
   getClassName() {
     return "GenAI_NanoGPT_v1";
   }
   inputPhase(e, i) {
-    return c(() => {
-      const s = this.wte.embed(e);
+    return l(() => {
+      const n = this.wte.embed(e);
       if (this.config.useRope === !1) {
-        const o = this.wpe.call(i, s);
+        const o = this.wpe.call(i, n);
         if (Array.isArray(o))
           throw new Error("PositionEmbedding output should not be an array");
         return o;
       }
-      return s;
+      return n;
     });
   }
-  forward(e, i, s) {
-    return this.validateInput(i), e.ropeCache = this.ropeCache, e.outputEmbeddings && (e.embeddings = []), c(() => {
+  forward(e, i, n) {
+    return this.validateInput(i), e.ropeCache = this.ropeCache, e.outputEmbeddings && (e.embeddings = []), l(() => {
       this.startMemory();
       let o = this.inputPhase(i, e);
       if (e.cache && e.cache.length !== this.blocks.length)
         throw console.error("Cache", e.cache), new Error(
           `Cache length ${e.cache.length} does not match number of blocks ${this.blocks.length}`
         );
-      for (let t = 0; t < this.blocks.length; t++) {
-        const r = this.blocks[t], d = Math.random() * 1e9, l = {
+      const t = e.mixedPrecision === !0 && _();
+      let s = t ? y(o) : o;
+      t && o !== s && o.dispose();
+      for (let r = 0; r < this.blocks.length; r++) {
+        const d = this.blocks[r], a = Math.random() * 1e9, m = {
           ...e,
-          seed: d,
-          pastKV: e.cache ? e.cache[t] : void 0
-        }, m = e.checkpointing && e.training ? r.callCheckpoint(l, o) : r.call(l, o);
-        e.outputEmbeddings ? (p(o), e.embeddings.push({ name: `block_output_${t}`, tensor: o })) : o.dispose(), o = m;
+          seed: a,
+          pastKV: e.cache ? e.cache[r] : void 0,
+          mixedPrecision: t
+        }, f = e.checkpointing && e.training ? d.callCheckpoint(m, s) : d.call(m, s);
+        e.outputEmbeddings ? (u(s), e.embeddings.push({ name: `block_output_${r}`, tensor: s })) : s.dispose(), s = f;
       }
-      o = this.lnF.call(e, o);
-      const n = this.wte.project(o);
-      e.outputEmbeddings ? (p(o), e.embeddings.push({ name: "final_norm_output", tensor: o })) : o.dispose();
-      let h;
-      return s && (h = this.calculateLoss(n, s)), this.endMemory("Forward"), h ? [n, h] : [n];
+      if (o = this.lnF.call({ ...e, mixedPrecision: t }, s), s.dispose(), e.skipLogits)
+        return this.endMemory("Forward"), [o];
+      const c = this.wte.project(o);
+      e.outputEmbeddings ? (u(o), e.embeddings.push({ name: "final_norm_output", tensor: o })) : o.dispose();
+      const h = t ? M(c) : c;
+      t && c !== h && c.dispose();
+      let p;
+      return n && (p = this.calculateLoss(h, n)), this.endMemory("Forward"), p ? [h, p] : [h];
     });
   }
   project(e) {
-    return c(() => this.wte.project(e));
+    return l(() => this.wte.project(e));
   }
   dispose() {
     this.wte.dispose(), this.wpe && this.wpe.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
   }
 }
 export {
-  R as default
+  I as default
 };

package/dist/models/model.d.ts

@@ -5,6 +5,7 @@ export interface ModelForwardAttributes extends ForwardAttributes {
     cache?: KVCache[];
     attentionScores?: AttentionScores;
     seed?: number;
+    skipLogits?: boolean;
 }
 interface TrainingState {
     steps: number;
@@ -13,6 +14,7 @@ interface TrainingState {
     loss: number;
 }
 export default abstract class Model<T extends ModelForwardAttributes> extends BaseLayer<T> {
+    lossScaling: number;
     trainingState: TrainingState | null;
     abstract getClassName(): string;
     abstract forward(attrs: T, idx: Tensor, targets?: Tensor): Tensor[];

package/dist/models/model.js

@@ -1,53 +1,60 @@
-import i from "../layers/BaseLayer.js";
-import "../index-BzFyqcy-.js";
+import m from "../layers/BaseLayer.js";
+import "../utilities/packed.js";
+import "../index-ZyQhjEPo.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 "../random_width-CXVRloNK.js";
-import "../register_all_kernels-DIGpEwcf.js";
-import "../index-Tf7vU29b.js";
-import "../dataset-DlZtKmBq.js";
+import "../random_width-DY6Kk2Dl.js";
+import "../register_all_kernels-Bwu1PTuU.js";
+import "../index-Cp39cXWe.js";
+import "../dataset-0xP8GjwI.js";
 import "../ops/cpu/rope.js";
 import "../ops/webgl/rope.js";
-import "../ops/grads/rope.js";
+import "../rope-B5UUMsPi.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/grads/softmax16.js";
+import "../matMul16--R5hOwDG.js";
+import "../ops/webgl/matMul16.js";
+import "../ops/cpu/matMul16.js";
+import "../pack16-CFUqumar.js";
+import "../ops/transpose16.js";
+import "../ops/reshape16.js";
+import "../ops/cpu/qkv.js";
+import "../ops/webgl/qkv.js";
+import "../ops/grads/qkv.js";
 import "../ops/cpu/normRMS.js";
 import "../ops/webgl/normRMS.js";
 import "../ops/grads/normRMS.js";
-import "../jszip.min-CjP2V1VV.js";
-import "../index-Dwqa6Zy2.js";
+import "../ops/grads/add16.js";
+import "../jszip.min-Bz5-11Bk.js";
+import "../index-DvYrXKkX.js";
 import "../Generator.js";
 import "../ops/cpu/adamAdjust.js";
 import "../ops/webgl/adamAdjust.js";
 import "../ops/cpu/adamMoments.js";
 import "../ops/webgl/adamMoments.js";
-import "../papaparse.min-C8l2Kvo1.js";
-import { estimateParameterCount as p } from "../utilities/parameters.js";
+import "../papaparse.min-C0cScC2i.js";
+import { estimateParameterCount as e } from "../utilities/parameters.js";
 import "../ops/cpu/scatterSub.js";
 import "../ops/webgl/scatterSub.js";
 import "../ops/cpu/gatherSub.js";
 import "../ops/webgl/gatherSub.js";
+import "../ops/cpu/matMulGelu.js";
+import "../ops/webgl/matMulGelu.js";
+import "../ops/grads/matMulGelu.js";
 import "../ops/cpu/gelu.js";
 import "../ops/webgl/gelu.js";
-import "../gelu-Bp_-935b.js";
+import "../gelu-CNLFZWea.js";
 import "../ops/webgl/log.js";
 import "../checks/normRMS.js";
 import "../checks/normRMSGrad.js";
-import { createSoftmaxCrossEntropyWithGrad as m } from "../training/sparseCrossEntropy.js";
-class x extends i {
+import { createSoftmaxCrossEntropyWithGrad as s } from "../training/sparseCrossEntropy.js";
+class st extends m {
+  lossScaling = 128;
   trainingState = null;
   getNumParams() {
-    return p(this.config);
+    return e(this.config);
   }
   validateInput(t) {
     if (t.shape.length !== 2)
@@ -57,14 +64,15 @@ class x extends i {
     if (t.dtype !== "int32")
       throw new Error(`Input tensor must be of type int32, got ${t.dtype}`);
   }
-  calculateLoss(t, o) {
+  calculateLoss(t, i) {
     try {
-      return m()(t, o).mean();
-    } catch (r) {
-      throw console.error("Error computing loss:", r), new Error(`Loss computation failed: ${r}`);
+      const r = s()(t, i), p = r.mean();
+      return r.dispose(), p;
+    } catch (o) {
+      throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
     }
   }
 }
 export {
-  x as default
+  st as default
 };

package/dist/{mulmat_packed_gpu-BT60jmzP.js → mulmat_packed_gpu-BMFhLwta.js}

@@ -1,20 +1,4 @@
-import { u as z } from "./gpgpu_math-CDaYiyE_.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
+import { u as z } from "./gpgpu_math-DDVJCn6-.js";
 class g {
   constructor(e, s, v, a = !1, r = !1, c = !1, t = null, o = !1, l = !1) {
     this.variableNames = ["matrixA", "matrixB"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = v, this.enableShapeUniforms = z(this.outputShape.length);

package/dist/{non_max_suppression_impl-CsEgBuMA.js → non_max_suppression_impl-B2W7YjZB.js}

@@ -1,19 +1,3 @@
-/**
- * @license
- * Copyright 2019 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
 function J(e, t, o) {
   const n = k(e, t, o), s = n < 0 ? -(n + 1) : n;
   e.splice(s, 0, t);
@@ -33,22 +17,6 @@ function A(e, t, o) {
   }
   return u ? n : -n - 1;
 }
-/**
- * @license
- * Copyright 2020 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
 function O(e, t, o, n, s) {
   return y(
     e,

package/dist/ones-CAMiP4I2.js

@@ -0,0 +1,15 @@
+import { E as a } from "./index-ZyQhjEPo.js";
+import { d as t, m as n, s as i } from "./tensor-DdQUJZlz.js";
+import { c as f } from "./complex-CSlYz-2T.js";
+import { z as c } from "./zeros-2gldETuK.js";
+function l(o, r = "float32") {
+  if (t(o), r === "complex64") {
+    const e = l(o, "float32"), m = c(o, "float32");
+    return f(e, m);
+  }
+  const s = n(i(o), r);
+  return a.makeTensor(s, o, r);
+}
+export {
+  l as o
+};

package/dist/ops/adamAdjust.js

@@ -1,4 +1,4 @@
-import { e as i } from "../index-BzFyqcy-.js";
+import { e as i } from "../index-ZyQhjEPo.js";
 import "./cpu/adamAdjust.js";
 import "./webgl/adamAdjust.js";
 function p(r, t, e, n, m, o) {

package/dist/ops/adamMoments.d.ts

@@ -1,2 +1,2 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-export declare function adamMoments(moments: Tensor, gradient: Tensor, beta1: number, beta2: number): Tensor;
+export declare function adamMoments(moments: Tensor, gradient: Tensor, beta1: number, beta2: number, lossScaling: number): Tensor;

package/dist/ops/adamMoments.js

@@ -1,9 +1,9 @@
-import { e as o } from "../index-BzFyqcy-.js";
+import { e as t } from "../index-ZyQhjEPo.js";
 import "./cpu/adamMoments.js";
 import "./webgl/adamMoments.js";
-function p(e, n, r, m) {
-  return o().runKernel("AdamMoments", { moments: e, gradient: n }, { beta1: r, beta2: m });
+function s(e, n, r, m, o) {
+  return t().runKernel("AdamMoments", { moments: e, gradient: n }, { beta1: r, beta2: m, lossScaling: o });
 }
 export {
-  p as adamMoments
+  s as adamMoments
 };

package/dist/ops/add16.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function add16(a: Tensor, b: Tensor): Tensor;

package/dist/ops/add16.js

@@ -0,0 +1,9 @@
+import { n as t, e as o } from "../index-ZyQhjEPo.js";
+import { isPackedTensor as n } from "../utilities/packed.js";
+import "./grads/add16.js";
+function m(r, e) {
+  return !n(r) && !n(e) ? t(r, e) : o().runKernel("Add16", { a: r, b: e });
+}
+export {
+  m as add16
+};

package/dist/ops/appendCache.js

@@ -1,15 +1,22 @@
-import { e as a } from "../index-BzFyqcy-.js";
+import { e as a } from "../index-ZyQhjEPo.js";
 import "./cpu/appendCache.js";
 import "./webgl/appendCache.js";
-import { c as s } from "../concat-B912vBbo.js";
-import { z as c } from "../zeros-Bj5rMYA7.js";
-function i(r, p, n, o) {
-  if (!o) {
-    const e = r.shape[2];
-    return s([r, c([r.shape[0], r.shape[1], p - e, r.shape[3]])], 2);
+import { isPackedTensor as c } from "../utilities/packed.js";
+import { c as t } from "../concat-BHlIJeyT.js";
+import { z as f } from "../zeros-2gldETuK.js";
+function C(r, o, n, p) {
+  if (!p) {
+    const e = r.shape[2], s = c(r);
+    return t(
+      [
+        r,
+        f([r.shape[0], r.shape[1], o - e, r.shape[3]], s ? "int32" : r.dtype)
+      ],
+      2
+    );
   }
-  return a().runKernel("AppendCache", { cache: o, item: r }, { maxSize: p, pastLen: n });
+  return a().runKernel("AppendCache", { cache: p, item: r }, { maxSize: o, pastLen: n });
 }
 export {
-  i as appendCache
+  C as appendCache
 };

package/dist/ops/attentionMask.js

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

package/dist/ops/concat16.d.ts

@@ -0,0 +1,2 @@
+import { Rank, Tensor } from '@tensorflow/tfjs-core';
+export declare function concat16<R extends Rank = Rank>(x: Tensor<R>[], axis?: number): Tensor<R>;

package/dist/ops/concat16.js

@@ -0,0 +1,9 @@
+import { isPackedTensor as o } from "../utilities/packed.js";
+import { e } from "../index-ZyQhjEPo.js";
+import { c } from "../concat-BHlIJeyT.js";
+function p(r, n) {
+  return o(r[0]) ? e().runKernel("Concat16", r, { axis: n ?? -1 }) : c(r, n);
+}
+export {
+  p as concat16
+};

package/dist/ops/cpu/adamAdjust.js

@@ -1,18 +1,19 @@
-import { f as k, o as t, q as i, a as q, w } from "../../index-BzFyqcy-.js";
-function z(a) {
-  const { moments: s, value: r } = a.inputs, { beta1: l, beta2: u, epsilon: m, learningRate: d } = a.attrs, e = s.shape.length, c = new Array(e).fill(0), n = s.shape.slice();
-  n[e - 1] = 1;
-  const o = c.slice();
-  o[e - 1] = 1;
-  const b = n.slice(), p = s.slice(c, n).squeeze([e - 1]), M = s.slice(o, b).squeeze([e - 1]), f = t(p, l), g = t(M, u);
-  return i(
-    q(t(f, i(w(g), m ?? 1e-8)), -d),
-    r
+import { l as t, n as r, m as k, o as z } from "../../index-ZyQhjEPo.js";
+import { r as A } from "../../tensor_util-DV-FP5Q3.js";
+function C(o) {
+  const { moments: n, value: i } = o.inputs, { beta1: l, beta2: m, epsilon: u, learningRate: d } = o.attrs, e = n.shape.length, c = new Array(e).fill(0), s = n.shape.slice();
+  s[e - 1] = 1;
+  const a = c.slice();
+  a[e - 1] = 1;
+  const p = s.slice(), b = n.slice(c, s).squeeze([e - 1]), M = n.slice(a, p).squeeze([e - 1]), f = t(b, l), g = t(M, m);
+  return r(
+    k(t(f, r(z(g), u ?? 1e-8)), -d),
+    i
   );
 }
-const A = {
+const h = {
   kernelName: "AdamAdjust",
   backendName: "cpu",
-  kernelFunc: z
+  kernelFunc: C
 };
-k(A);
+A(h);

package/dist/ops/cpu/adamMoments.js

@@ -1,16 +1,17 @@
-import { f as p } from "../../index-BzFyqcy-.js";
-import { s as f } from "../../stack-DFatutCx.js";
-function b(t) {
-  const { moments: n, gradient: c } = t.inputs, { beta1: o, beta2: m } = t.attrs, e = n.shape.length, a = new Array(e).fill(0), s = n.shape.slice();
+import "../../index-ZyQhjEPo.js";
+import { r as p } from "../../tensor_util-DV-FP5Q3.js";
+import { s as b } from "../../stack-yOIAalTq.js";
+function f(t) {
+  const { moments: n, gradient: o } = t.inputs, { beta1: c, beta2: m } = t.attrs, e = n.shape.length, a = new Array(e).fill(0), s = n.shape.slice();
   s[e - 1] = 1;
-  const i = a.slice();
-  i[e - 1] = 1;
-  const r = s.slice(), l = n.slice(a, s).squeeze([e - 1]), u = n.slice(i, r).squeeze([e - 1]), M = l.mul(o).add(c.mul(1 - o)), d = u.mul(m).add(c.square().mul(1 - m));
-  return f([M, d], -1);
+  const r = a.slice();
+  r[e - 1] = 1;
+  const i = s.slice(), l = n.slice(a, s).squeeze([e - 1]), u = n.slice(r, i).squeeze([e - 1]), M = l.mul(c).add(o.mul(1 - c)), d = u.mul(m).add(o.square().mul(1 - m));
+  return b([M, d], -1);
 }
 const g = {
   kernelName: "AdamMoments",
   backendName: "cpu",
-  kernelFunc: b
+  kernelFunc: f
 };
 p(g);

package/dist/ops/cpu/appendCache.js

@@ -1,13 +1,14 @@
-import { f as d } from "../../index-BzFyqcy-.js";
-import { c as h } from "../../concat-B912vBbo.js";
+import "../../index-ZyQhjEPo.js";
+import { r as d } from "../../tensor_util-DV-FP5Q3.js";
+import { c as h } from "../../concat-BHlIJeyT.js";
 function u(p) {
-  const { cache: n, item: s } = p.inputs, { maxSize: i, pastLen: c } = p.attrs, t = n.shape[0], o = n.shape[1], a = n.shape[3], e = s.shape[2];
-  if (c + e <= i) {
-    const l = n.slice([0, 0, 0, 0], [t, o, c, a]), m = n.slice([0, 0, c + e, 0], [t, o, i - c - e, a]), r = e < e ? s.slice([0, 0, 0, 0], [t, o, e, a]) : s, k = h([l, r, m], 2);
-    return l.dispose(), m.dispose(), r !== s && r.dispose(), k;
+  const { cache: n, item: s } = p.inputs, { maxSize: a, pastLen: c } = p.attrs, t = n.shape[0], o = n.shape[1], r = n.shape[3], e = s.shape[2];
+  if (c + e <= a) {
+    const m = n.slice([0, 0, 0, 0], [t, o, c, r]), f = n.slice([0, 0, c + e, 0], [t, o, a - c - e, r]), i = e < e ? s.slice([0, 0, 0, 0], [t, o, e, r]) : s, k = h([m, i, f], 2);
+    return m.dispose(), f.dispose(), i !== s && i.dispose(), k;
   }
-  const f = n.slice([0, 0, e, 0], [t, o, i - e, a]), C = h([f, s], 2);
-  return f.dispose(), C;
+  const l = n.slice([0, 0, e, 0], [t, o, a - e, r]), C = h([l, s], 2);
+  return l.dispose(), C;
 }
 const w = {
   kernelName: "AppendCache",

package/dist/ops/cpu/attentionMask.js

@@ -1,21 +1,22 @@
-import { f as a, h as p, b as u } from "../../index-BzFyqcy-.js";
-import { l as N, w as b } from "../../ops-LuCMAnmM.js";
-import { o as A } from "../../ones-tIJeHlq-.js";
-import { z as I } from "../../zeros-Bj5rMYA7.js";
-import { m as g } from "../../mat_mul-DzjTFx-u.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(A([t, t]), -1, 0).cast("bool"), l = I([t, t]), i = p([t, t], Number.NEGATIVE_INFINITY), f = b(m, l, i), k = g(s, e, !1, !0).mul(u(r)), d = f.slice([0, 0], [c, t]).expandDims(0).expandDims(0);
-  return k.add(d);
+import { i as d, b as u } from "../../index-ZyQhjEPo.js";
+import { r as o } from "../../tensor_util-DV-FP5Q3.js";
+import { l as N, w as b } from "../../ops-CNI3TwqM.js";
+import { o as A } from "../../ones-CAMiP4I2.js";
+import { z as I } from "../../zeros-2gldETuK.js";
+import { m as g } from "../../mat_mul-DeAh4uTH.js";
+function a(n) {
+  const { q: s, k: e } = n.inputs, { divisor: r } = n.attrs, c = s.shape[2], t = e.shape[2], m = N.bandPart(A([t, t]), -1, 0).cast("bool"), i = I([t, t]), l = d([t, t], Number.NEGATIVE_INFINITY), f = b(m, i, l), k = g(s, e, !1, !0).mul(u(r)), p = f.slice([0, 0], [c, t]).expandDims(0).expandDims(0);
+  return k.add(p);
 }
-const h = {
+const w = {
   kernelName: "AttentionMask",
   backendName: "cpu",
-  kernelFunc: o
+  kernelFunc: a
 };
-a(h);
-const w = {
+o(w);
+const M = {
   kernelName: "AttentionMask",
   backendName: "tensorflow",
-  kernelFunc: o
+  kernelFunc: a
 };
-a(w);
+o(M);

package/dist/ops/cpu/fusedSoftmax.js

@@ -1,29 +1,30 @@
-import { f as e } from "../../index-BzFyqcy-.js";
-import { s as f } from "../../softmax-D7Jj3p_P.js";
-function n(t) {
-  const { inputs: s, attrs: a } = t, { logits: o } = s, { dim: i, dropoutRate: r } = a;
-  if (!o)
+import "../../index-ZyQhjEPo.js";
+import { r as e } from "../../tensor_util-DV-FP5Q3.js";
+import { s as m } from "../../softmax-ZHVebtR1.js";
+function o(t) {
+  const { inputs: s, attrs: a } = t, { logits: n } = s, { dim: i, dropoutRate: r } = a;
+  if (!n)
     throw new Error("Error in softmax: input logits is null");
-  return r !== void 0 && r > 0 && console.warn("Dropout in fusedSoftmax not implemented for CPU backend, skipping dropout."), f(o, i);
+  return r !== void 0 && r > 0 && console.warn("Dropout in fusedSoftmax not implemented for CPU backend, skipping dropout."), m(n, i);
 }
-const m = {
+const f = {
   kernelName: "FusedSoftmax",
   backendName: "cpu",
-  kernelFunc: n
+  kernelFunc: o
 };
-e(m);
+e(f);
 const u = {
   kernelName: "FusedSoftmax",
   backendName: "tensorflow",
-  kernelFunc: n
+  kernelFunc: o
 };
 e(u);
 const l = {
   kernelName: "FusedSoftmax",
   backendName: "webgpu",
-  kernelFunc: n
+  kernelFunc: o
 };
 e(l);
 export {
-  n as softmaxCPU
+  o as softmaxCPU
 };

package/dist/ops/cpu/gatherSub.js

@@ -1,34 +1,19 @@
-import { E as u, F as c, G as g, a8 as h, f as m, c as p } from "../../index-BzFyqcy-.js";
-import { r as f } from "../../range-CWcz7xFA.js";
-import { s as l } from "../../stack-DFatutCx.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
+import { A as u, B as c, E as m, c as g } from "../../index-ZyQhjEPo.js";
+import { k as p, r as h } from "../../tensor_util-DV-FP5Q3.js";
+import { r as f } from "../../range-BMS52eQi.js";
+import { s as l } from "../../stack-yOIAalTq.js";
 function N(e, t) {
-  const n = c(t, "indices", "gatherND", "int32"), s = { params: c(e, "x", "gatherND", "string_or_numeric"), indices: n };
-  return g.runKernel(h, s);
+  const n = c(t, "indices", "gatherND", "int32"), r = { params: c(e, "x", "gatherND", "string_or_numeric"), indices: n };
+  return m.runKernel(p, r);
 }
 const b = /* @__PURE__ */ u({ gatherND_: N });
 function d(e) {
-  const { values: t, labels: n, logits: r } = e.inputs, s = n.shape[0], a = f(0, s, 1, "int32"), i = l([a, n], 1), o = b(r, i);
-  return p(t, o);
+  const { values: t, labels: n, logits: s } = e.inputs, r = n.shape[0], o = f(0, r, 1, "int32"), i = l([o, n], 1), a = b(s, i);
+  return g(t, a);
 }
 const k = {
   kernelName: "EfficientGatherSub",
   backendName: "cpu",
   kernelFunc: d
 };
-m(k);
+h(k);