@genai-fi/nanogpt 0.7.0 → 0.7.2

package/dist/ops/gelu.js

@@ -1,7 +1,7 @@
-import "../index-C0dhsYom.js";
+import "../index-BoWRt-10.js";
 import "./cpu/gelu.js";
 import "./webgl/gelu.js";
-import { d as e, g as i } from "../gelu-uyHP1x1f.js";
+import { d as e, g as i } from "../gelu-C-dPj6Ku.js";
 export {
   e as dGelu,
   i as gelu

package/dist/ops/grads/attentionMask.js

@@ -1,4 +1,4 @@
-import { i as m, b as i } from "../../index-C0dhsYom.js";
+import { i as m, b as i } from "../../index-BoWRt-10.js";
 import { matMulMul as u } from "../matMulMul.js";
 const p = {
   kernelName: "AttentionMask",

package/dist/ops/grads/fusedSoftmax.js

@@ -1,6 +1,6 @@
-import { i as f, a as i, c as l } from "../../index-C0dhsYom.js";
+import { i as f, a as i, c as l } from "../../index-BoWRt-10.js";
 import { mulDrop as g } from "../mulDrop.js";
-import { s as T } from "../../sum-B8wEpKsg.js";
+import { s as T } from "../../sum-66ew2byf.js";
 const Y = {
   kernelName: "FusedSoftmax",
   outputsToSave: [!0],

package/dist/ops/grads/gelu.js

@@ -1,5 +1,5 @@
-import "../../index-C0dhsYom.js";
-import { a as e } from "../../gelu-uyHP1x1f.js";
+import "../../index-BoWRt-10.js";
+import { a as e } from "../../gelu-C-dPj6Ku.js";
 export {
   e as geluGradConfig
 };

package/dist/ops/grads/matMulGelu.js

@@ -1,4 +1,4 @@
-import { i as a, e as o } from "../../index-C0dhsYom.js";
+import { i as a, e as o } from "../../index-BoWRt-10.js";
 function i(e, n, r) {
   return o().runKernel("MatMulGeluGrad", { dy: e, x: n, kernel: r });
 }

package/dist/ops/grads/normRMS.js

@@ -1,4 +1,4 @@
-import { i as t, e as g } from "../../index-C0dhsYom.js";
+import { i as t, e as g } from "../../index-BoWRt-10.js";
 function i(r, a, n) {
   return g().runKernel("RMSNormGrad", { dy: r, x: a, gamma: n });
 }

package/dist/ops/grads/qkv.js

@@ -1,4 +1,4 @@
-import { i as Q } from "../../index-C0dhsYom.js";
+import { i as Q } from "../../index-BoWRt-10.js";
 const V = {
   kernelName: "QKV",
   inputsToSave: ["x", "kernel"],

package/dist/ops/grads/rope.js

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

package/dist/ops/matMulGelu.js

@@ -1,4 +1,4 @@
-import { e as u } from "../index-C0dhsYom.js";
+import { e as u } from "../index-BoWRt-10.js";
 import "./cpu/matMulGelu.js";
 import "./webgl/matMulGelu.js";
 import "./grads/matMulGelu.js";

package/dist/ops/matMulMul.js

@@ -1,4 +1,4 @@
-import { e as u } from "../index-C0dhsYom.js";
+import { e as u } from "../index-BoWRt-10.js";
 import "./cpu/matMulMul.js";
 import "./webgl/matMulMul.js";
 function m(e, r, t, l = !1, n = !1) {

package/dist/ops/mulDrop.js

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

package/dist/ops/normRMS.js

@@ -1,4 +1,4 @@
-import { e as n } from "../index-C0dhsYom.js";
+import { e as n } from "../index-BoWRt-10.js";
 import "./cpu/normRMS.js";
 import "./webgl/normRMS.js";
 import "./grads/normRMS.js";

package/dist/ops/qkv.js

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

package/dist/ops/rope.js

@@ -1,8 +1,8 @@
-import { e as p } from "../index-C0dhsYom.js";
-import "../random_width-DkYP8W8N.js";
-import "../register_all_kernels-guvSxp7M.js";
+import { e as p } from "../index-BoWRt-10.js";
+import "../random_width-sZORGo5k.js";
+import "../register_all_kernels-BwDSRN-f.js";
 import "../index-Tf7vU29b.js";
-import "../dataset-JDyjG3QR.js";
+import "../dataset-CtdBYwjo.js";
 import "./cpu/rope.js";
 import "./webgl/rope.js";
 import "./grads/rope.js";

package/dist/ops/scatterSub.js

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

package/dist/ops/webgl/adamAdjust.d.ts

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

package/dist/ops/webgl/adamAdjust.js

@@ -0,0 +1,50 @@
+import { r as n } from "../../Reshape-DH5srBP0.js";
+import { f } from "../../index-BoWRt-10.js";
+class v {
+  variableNames = ["moments", "value"];
+  outputShape;
+  userCode;
+  customUniforms = [
+    { name: "invBeta1", type: "float" },
+    { name: "invBeta2", type: "float" },
+    { name: "learningRate", type: "float" },
+    { name: "epsilon", type: "float" }
+  ];
+  constructor(a) {
+    this.outputShape = a, this.userCode = `
+        void main() {
+            float v = getValueAtOutCoords();
+            int coords = getOutputCoords();
+            coords *= 2;
+            float m1 = getMoments(coords);
+            float m2 = getMoments(coords + 1);
+
+            float m1Hat = m1 * invBeta1;
+            float m2Hat = m2 * invBeta2;
+
+            float invSqrt = inversesqrt(max(m2Hat, 1e-30));
+            float invDenom = invSqrt / (1.0 + epsilon * invSqrt);
+            float adjustedValue = -learningRate * m1Hat * invDenom + v;
+
+            setOutput(adjustedValue);
+        }
+        `;
+  }
+}
+function h(t) {
+  const { moments: a, value: s } = t.inputs, { beta1: i, beta2: l, learningRate: u, epsilon: p } = t.attrs, e = t.backend, r = n({ inputs: { x: a }, backend: e, attrs: { shape: [-1] } }), o = n({ inputs: { x: s }, backend: e, attrs: { shape: [-1] } }), d = new v(o.shape), m = e.runWebGLProgram(d, [r, o], "float32", [
+    [1 / i],
+    [1 / l],
+    [u],
+    [p]
+  ]);
+  e.disposeIntermediateTensorInfo(r), e.disposeIntermediateTensorInfo(o);
+  const c = n({ inputs: { x: m }, backend: e, attrs: { shape: s.shape } });
+  return e.disposeIntermediateTensorInfo(m), c;
+}
+const g = {
+  kernelName: "AdamAdjust",
+  backendName: "webgl",
+  kernelFunc: h
+};
+f(g);

package/dist/ops/webgl/adamMoments.d.ts

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

package/dist/ops/webgl/adamMoments.js

@@ -0,0 +1,40 @@
+import { f as m } from "../../index-BoWRt-10.js";
+class i {
+  variableNames = ["moments", "gradient"];
+  outputShape;
+  userCode;
+  customUniforms = [
+    { name: "beta1", type: "float" },
+    { name: "beta2", type: "float" }
+  ];
+  constructor(e) {
+    this.outputShape = e;
+    const t = e.length, n = t === 1 ? "int" : `ivec${Math.min(t, 4)}`, a = t === 1 ? "coords" : `coords[${t - 1}]`, r = t === 1 ? "coords" : e.slice(0, -1).map((s, c) => `coords[${c}]`).join(", ");
+    this.userCode = `
+        void main() {
+            float m = getMomentsAtOutCoords();
+            ${n} coords = getOutputCoords();
+
+            // Add gradient clipping here
+            float g = clamp(getGradient(${r}), -1.0, 1.0);
+            int which = ${a};
+
+            float beta = which == 0 ? beta1 : beta2;
+            float gg = which == 0 ? g : g * g;
+    
+            float newM = m * beta + gg * (1.0 - beta);
+            setOutput(newM);
+        }
+        `;
+  }
+}
+function d(o) {
+  const { moments: e, gradient: t } = o.inputs, { beta1: n, beta2: a } = o.attrs, r = o.backend, s = new i(e.shape);
+  return r.runWebGLProgram(s, [e, t], "float32", [[n], [a]]);
+}
+const g = {
+  kernelName: "AdamMoments",
+  backendName: "webgl",
+  kernelFunc: d
+};
+m(g);

package/dist/ops/webgl/appendCache.js

@@ -1,4 +1,4 @@
-import { f as p } from "../../index-C0dhsYom.js";
+import { f as p } from "../../index-BoWRt-10.js";
 class m {
   variableNames = ["cache", "item"];
   outputShape;

package/dist/ops/webgl/attentionMask.js

@@ -1,4 +1,4 @@
-import { f as m } from "../../index-C0dhsYom.js";
+import { f as m } from "../../index-BoWRt-10.js";
 class h {
   variableNames = ["q", "k"];
   outputShape;

package/dist/ops/webgl/fusedSoftmax.js

@@ -1,7 +1,7 @@
-import { m as b, s as I, r as k } from "../../RealDiv-C4hOvYOZ.js";
-import { r as v } from "../../Reshape-BLijOA8h.js";
-import { f as w, p as P } from "../../index-C0dhsYom.js";
-import { e as S } from "../../axis_util-DaAl5MER.js";
+import { m as b, s as I, r as k } from "../../RealDiv-Dy0p8Bvo.js";
+import { r as v } from "../../Reshape-DH5srBP0.js";
+import { f as w, p as P } from "../../index-BoWRt-10.js";
+import { e as S } from "../../axis_util-BzbKo31C.js";
 class T {
   variableNames = ["logits", "maxLogits"];
   outputShape;

package/dist/ops/webgl/gatherSub.js

@@ -1,17 +1,17 @@
-import { f as l } from "../../index-C0dhsYom.js";
+import { f as l } from "../../index-BoWRt-10.js";
 class u {
   variableNames = ["labels", "logits", "values"];
   outputShape;
   userCode;
   constructor(e) {
     this.outputShape = [e], this.userCode = `
-      void main() {
-        int coords = getOutputCoords();
-        int index = int(getLabelsAtOutCoords());
-        float val = getValuesAtOutCoords();
-        float logit = getLogits(coords, index);
-        setOutput(val - logit);
-      }
+          void main() {
+              int coords = getOutputCoords();
+              int index = int(getLabelsAtOutCoords());
+              float val = getValuesAtOutCoords();
+              float logit = getLogits(coords, index);
+              setOutput(val - logit);
+          }
     `;
   }
 }

package/dist/ops/webgl/gelu.js

@@ -1,5 +1,5 @@
-import { f as a } from "../../index-C0dhsYom.js";
-import { u as s, C as x } from "../../kernel_funcs_utils-CwRTFqrc.js";
+import { f as a } from "../../index-BoWRt-10.js";
+import { u as s, C as x } from "../../kernel_funcs_utils-BYKWV8Aa.js";
 const t = 0.7978845608028654, r = 0.044715, c = x + `
     float x3 = x * x * x;
     float inner = x + ${r} * x3;

package/dist/ops/webgl/log.js

@@ -1,6 +1,6 @@
-import { f as e, a8 as r } from "../../index-C0dhsYom.js";
-import { u as s, l as N } from "../../kernel_funcs_utils-CwRTFqrc.js";
-import { y as l } from "../../shared-BNa2q6jD.js";
+import { f as e, af as r } from "../../index-BoWRt-10.js";
+import { u as s, l as N } from "../../kernel_funcs_utils-BYKWV8Aa.js";
+import { y as l } from "../../shared-wS99K7_n.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.

package/dist/ops/webgl/matMulGelu.js

@@ -1,7 +1,7 @@
-import { f as _, t as R, e as C, j as A, k as N, l as H, u as O } from "../../index-C0dhsYom.js";
-import { r as f } from "../../Reshape-BLijOA8h.js";
-import { M as U } from "../../mulmat_packed_gpu-Gzf3I9UV.js";
-import { m as E } from "../../mat_mul-1nwdPkQ_.js";
+import { f as _, t as R, e as C, j as A, k as N, l as H, u as O } from "../../index-BoWRt-10.js";
+import { r as f } from "../../Reshape-DH5srBP0.js";
+import { M as U } from "../../mulmat_packed_gpu-VSekgsNv.js";
+import { m as E } from "../../mat_mul-8m8pfdcx.js";
 const M = 0.7978845608028654, x = 0.044715, j = `
     vec4 x3 = x * x * x;
     vec4 inner = x + ${x} * x3;

package/dist/ops/webgl/matMulMul.js

@@ -1,4 +1,4 @@
-import { f as u } from "../../index-C0dhsYom.js";
+import { f as u } from "../../index-BoWRt-10.js";
 import { batchMatMulGeluImpl as c } from "./matMulGelu.js";
 const M = `
     return a * b;

package/dist/ops/webgl/mulDropout.js

@@ -1,4 +1,4 @@
-import { f as m } from "../../index-C0dhsYom.js";
+import { f as m } from "../../index-BoWRt-10.js";
 class f {
   variableNames = ["a", "b"];
   outputShape;

package/dist/ops/webgl/normRMS.js

@@ -1,5 +1,5 @@
-import { f as p, e as G } from "../../index-C0dhsYom.js";
-import { s as x } from "../../sum-B8wEpKsg.js";
+import { f as p, e as G } from "../../index-BoWRt-10.js";
+import { s as x } from "../../sum-66ew2byf.js";
 class y {
   variableNames = ["x", "meanSquare", "gamma"];
   outputShape;

package/dist/ops/webgl/qkv.js

@@ -1,4 +1,4 @@
-import { f as i } from "../../index-C0dhsYom.js";
+import { f as i } from "../../index-BoWRt-10.js";
 class l {
   variableNames = ["x", "kernel"];
   outputShape;

package/dist/ops/webgl/rope.js

@@ -1,4 +1,4 @@
-import { f as u } from "../../index-C0dhsYom.js";
+import { f as u } from "../../index-BoWRt-10.js";
 class l {
   variableNames = ["x", "sin", "cos"];
   outputShape;

package/dist/ops/webgl/scatterSub.js

@@ -1,4 +1,4 @@
-import { f as i } from "../../index-C0dhsYom.js";
+import { f as i } from "../../index-BoWRt-10.js";
 class u {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;

package/dist/ops/webgpu/adamAdjust.d.ts

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

package/dist/ops/webgpu/adamAdjust.js

@@ -0,0 +1,54 @@
+import { g as p } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as d, c as l } from "../../webgpu_util-D____QpY.js";
+import { f, a2 as c } from "../../index-BoWRt-10.js";
+class h {
+  variableNames = ["moments", "value"];
+  outputShape;
+  shaderKey = "AdamAdjust";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "invbeta1: f32, invbeta2: f32, learningRate: f32, epsilon: f32";
+  outputComponent = 1;
+  variableComponents = [2, 1];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = d(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+        ${p("index")} {
+            if (index < uniforms.size) {
+                let moments: vec2<f32> = moments[index];
+                let value: f32 = value[index];
+
+                let m1Hat = moments.x * uniforms.invbeta1;
+                let m2Hat = moments.y * uniforms.invbeta2;
+
+                let invSqrt = inverseSqrt(max(m2Hat, 1e-30));
+                let invDenom = invSqrt / fma(uniforms.epsilon, invSqrt, 1.0);
+                let adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);
+
+                setOutputAtIndex(index, adjustedValue);
+            }
+        }
+    `;
+  }
+}
+function v(t) {
+  const { moments: e, value: a } = t.inputs, { beta1: n, beta2: s, learningRate: o, epsilon: i } = t.attrs, r = t.backend;
+  c(e.shape, [...a.shape, 2], "Error in AdamAdjust: ");
+  const u = new h(a.shape), m = [
+    { type: "float32", data: [1 / n] },
+    { type: "float32", data: [1 / s] },
+    { type: "float32", data: [o] },
+    { type: "float32", data: [i] }
+  ];
+  return r.runWebGPUProgram(u, [e, a], "float32", m);
+}
+const b = {
+  kernelName: "AdamAdjust",
+  backendName: "webgpu",
+  kernelFunc: v
+};
+f(b);

package/dist/ops/webgpu/adamMoments.d.ts

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

package/dist/ops/webgpu/adamMoments.js

@@ -0,0 +1,58 @@
+import { g as m } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as u, c as p } from "../../webgpu_util-D____QpY.js";
+import { f as d, a2 as c } from "../../index-BoWRt-10.js";
+class h {
+  variableNames = ["moments", "gradient"];
+  outputShape;
+  shaderKey = "AdamMoments";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "beta1: f32, beta2: f32";
+  outputComponent = 2;
+  variableComponents = [2, 1];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = u(this.outputShape.slice(0, -1)), this.dispatch = p(
+      this.dispatchLayout,
+      this.outputShape.slice(0, -1),
+      this.workgroupSize,
+      [1, 1, 1]
+    );
+  }
+  getUserCode() {
+    return `
+        ${m("index")} {
+            if (index < uniforms.size) {
+                let m: vec2<f32> = moments[index];
+
+                // Add gradient clipping here
+                let g: f32 = clamp(gradient[index], -1.0, 1.0);
+
+                let newM1 = fma(m.x, uniforms.beta1, g * (1.0 - uniforms.beta1));
+                let newM2 = fma(m.y, uniforms.beta2, g * g * (1.0 - uniforms.beta2));
+
+                setOutputAtIndex(index, vec2<f32>(newM1, newM2));
+            }
+        }
+    `;
+  }
+}
+function f(t) {
+  const { moments: e, gradient: o } = t.inputs, { beta1: a, beta2: n } = t.attrs, r = t.backend;
+  if (c(e.shape, [...o.shape, 2], "Error in AdamMoments: "), a < 0 || a >= 1)
+    throw new Error(`Invalid beta1 value: ${a}. Must be in the range [0, 1).`);
+  if (n < 0 || n >= 1)
+    throw new Error(`Invalid beta2 value: ${n}. Must be in the range [0, 1).`);
+  const s = new h(e.shape), i = [
+    { type: "float32", data: [a] },
+    { type: "float32", data: [n] }
+  ];
+  return r.runWebGPUProgram(s, [e, o], "float32", i);
+}
+const l = {
+  kernelName: "AdamMoments",
+  backendName: "webgpu",
+  kernelFunc: f
+};
+d(l);

package/dist/ops/webgpu/appendCache.js

@@ -1,6 +1,7 @@
-import { f as u, c as d, g as l } from "../../webgpu_util-g13LvDIv.js";
-import { f as m } from "../../index-C0dhsYom.js";
-class f {
+import { g as u } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as d, c as m } from "../../webgpu_util-D____QpY.js";
+import { f as l, a2 as f } from "../../index-BoWRt-10.js";
+class g {
   variableNames = ["cache", "item"];
   outputShape;
   shaderKey = "AppendCache";
@@ -9,14 +10,14 @@ class f {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "cacheT: i32";
-  constructor(t, a, s, o, c) {
-    const i = Math.min(s + 1, c);
-    this.outputShape = [t, a, i, o], this.dispatchLayout = u(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(e, t, s, a, c) {
+    const o = Math.min(s + 1, c);
+    this.shaderKey = `AppendCache_${o}`, this.outputShape = [e, t, o, a], this.dispatchLayout = d(this.outputShape), this.dispatch = m(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
-    const t = this.outputShape[2];
+    const e = this.outputShape[2];
     return `
-        ${l("index")} {
+        ${u("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -25,7 +26,7 @@ class f {
                 let d = coords[3];
 
                 let itemT = 1;
-                let maxSize = ${t};
+                let maxSize = ${e};
                 let totalT = uniforms.cacheT + itemT;
                 let start = select(0, 1, totalT >= maxSize);
 
@@ -33,24 +34,27 @@ class f {
                 var val = 0.0;
                 if (srcT < uniforms.cacheT) {
                     val = getCache(b, h, srcT, d);
-                } else if (srcT == uniforms.cacheT) {
+                }
+                if (srcT == uniforms.cacheT) {
                     val = getItem(b, h, 0, d);
-                } else {
-                    val = 0.0;
                 }
+
                 setOutputAtIndex(index, val);
             }
         }
         `;
   }
 }
-function T(e) {
-  const { cache: t, item: a } = e.inputs, { maxSize: s, pastLen: o } = e.attrs, c = e.backend, i = t.shape[0], r = t.shape[2], n = t.shape[1], h = new f(i, n, r, a.shape[3], s), p = [{ type: "int32", data: [o] }];
-  return c.runWebGPUProgram(h, [t, a], "float32", p);
+function S(r) {
+  const { cache: e, item: t } = r.inputs, { maxSize: s, pastLen: a } = r.attrs, c = r.backend, o = e.shape[0], n = e.shape[2], i = e.shape[1];
+  if (f(t.shape, [o, i, 1, t.shape[3]], "Error in AppendCache: "), a < 0 || a > s)
+    throw new Error(`Invalid pastLen value: ${a}. Must be in the range [0, ${s}].`);
+  const h = new g(o, i, n, t.shape[3], s), p = [{ type: "int32", data: [a] }];
+  return c.runWebGPUProgram(h, [e, t], "float32", p);
 }
-const g = {
+const T = {
   kernelName: "AppendCache",
   backendName: "webgpu",
-  kernelFunc: T
+  kernelFunc: S
 };
-m(g);
+l(T);

package/dist/ops/webgpu/attentionMask.js

@@ -1,5 +1,6 @@
-import { f } from "../../index-C0dhsYom.js";
-import { f as m, c as k, g as l } from "../../webgpu_util-g13LvDIv.js";
+import { f, a2 as m } from "../../index-BoWRt-10.js";
+import { g as k } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as l, c as v } from "../../webgpu_util-D____QpY.js";
 class g {
   variableNames = ["q", "k"];
   outputShape;
@@ -13,13 +14,13 @@ class g {
   nh;
   T1;
   T2;
-  constructor(t, e, o, i, a) {
-    if (this.outputShape = [t, e, o, i], this.hs = a, this.nh = e, this.T1 = o, this.T2 = i, this.dispatchLayout = m(this.outputShape), this.dispatch = k(this.dispatchLayout, this.outputShape, this.workgroupSize), a % 4 !== 0)
+  constructor(t, e, s, o, i) {
+    if (this.shaderKey = `AttentionMask_${i}`, this.outputShape = [t, e, s, o], this.hs = i, this.nh = e, this.T1 = s, this.T2 = o, this.dispatchLayout = l(this.outputShape), this.dispatch = v(this.dispatchLayout, this.outputShape, this.workgroupSize), i % 4 !== 0)
       throw new Error("Head size must be a multiple of 4 for AttentionMaskProgram");
   }
   getUserCode() {
     return `
-            ${l("index")} {
+            ${k("index")} {
                 
                 let coords = getCoordsFromIndex(index);
                 let b = coords[0];
@@ -32,13 +33,14 @@ class g {
                         setOutputAtIndex(index, uniforms.inf);
                         return;
                     }
+
+                    let q0 = getIndexFromCoords4D(vec4<i32>(b, h, t1, 0), uniforms.qShape);
+                    let k0 = getIndexFromCoords4D(vec4<i32>(b, h, t2, 0), uniforms.kShape);
                     
                     var sum: f32 = 0.0;
                     for (var i: i32 = 0; i < ${this.hs}; i = i + 4) {
-                        let q0 = getIndexFromCoords4D(vec4<i32>(b, h, t1, i), uniforms.qShape);
-                        let qv = vec4<f32>(q[q0], q[q0 + 1], q[q0 + 2], q[q0 + 3]);
-                        let k0 = getIndexFromCoords4D(vec4<i32>(b, h, t2, i), uniforms.kShape);
-                        let kv = vec4<f32>(k[k0], k[k0 + 1], k[k0 + 2], k[k0 + 3]);
+                        let qv = vec4<f32>(q[q0 + i], q[q0 + i + 1], q[q0 + i + 2], q[q0 + i + 3]);
+                        let kv = vec4<f32>(k[k0 + i], k[k0 + i + 1], k[k0 + i + 2], k[k0 + i + 3]);
                         sum = sum + dot(qv, kv);
                     }
                     let scaled = sum * uniforms.divisor;
@@ -48,17 +50,22 @@ class g {
         `;
   }
 }
-function q(s) {
-  const { q: t, k: e } = s.inputs, { divisor: o, pastLen: i } = s.attrs, a = s.backend, n = t.shape[0], r = t.shape[2], u = e.shape[2], c = t.shape[1], d = t.shape[3], h = new g(n, c, r, u, d), p = [
-    { type: "float32", data: [o] },
-    { type: "int32", data: [i] },
+function b(n) {
+  const { q: t, k: e } = n.inputs, { divisor: s, pastLen: o } = n.attrs, i = n.backend, r = t.shape[0], p = t.shape[2], a = e.shape[2], u = t.shape[1], h = t.shape[3];
+  if (m(e.shape, [r, u, a, h], "Error in AttentionMask: "), s === 0)
+    throw new Error("Divisor must be non-zero in AttentionMask");
+  if (o < 0)
+    throw new Error("pastLen must be non-negative in AttentionMask");
+  const c = new g(r, u, p, a, h), d = [
+    { type: "float32", data: [s] },
+    { type: "int32", data: [o] },
     { type: "float32", data: [Number.NEGATIVE_INFINITY] }
   ];
-  return a.runWebGPUProgram(h, [t, e], "float32", p);
+  return i.runWebGPUProgram(c, [t, e], "float32", d);
 }
-const v = {
+const q = {
   kernelName: "AttentionMask",
   backendName: "webgpu",
-  kernelFunc: q
+  kernelFunc: b
 };
-f(v);
+f(q);