@genai-fi/nanogpt 0.7.1 → 0.7.3

package/dist/ops/scatterSub.js

@@ -1,4 +1,4 @@
-import { e as i } from "../index-UdZhlibC.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.js

@@ -1,5 +1,5 @@
-import { r as n } from "../../Reshape-CEsEp0AI.js";
-import { f } from "../../index-UdZhlibC.js";
+import { r as n } from "../../Reshape-DH5srBP0.js";
+import { f } from "../../index-BoWRt-10.js";
 class v {
   variableNames = ["moments", "value"];
   outputShape;

package/dist/ops/webgl/adamMoments.js

@@ -1,4 +1,4 @@
-import { f as m } from "../../index-UdZhlibC.js";
+import { f as m } from "../../index-BoWRt-10.js";
 class i {
   variableNames = ["moments", "gradient"];
   outputShape;
@@ -9,12 +9,14 @@ class i {
   ];
   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}]`, s = t === 1 ? "coords" : e.slice(0, -1).map((r, c) => `coords[${c}]`).join(", ");
+    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();
-            float g = getGradient(${s});
+
+            // Add gradient clipping here
+            float g = clamp(getGradient(${r}), -1.0, 1.0);
             int which = ${a};
 
             float beta = which == 0 ? beta1 : beta2;
@@ -27,8 +29,8 @@ class i {
   }
 }
 function d(o) {
-  const { moments: e, gradient: t } = o.inputs, { beta1: n, beta2: a } = o.attrs, s = o.backend, r = new i(e.shape);
-  return s.runWebGLProgram(r, [e, t], "float32", [[n], [a]]);
+  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",

package/dist/ops/webgl/appendCache.js

@@ -1,4 +1,4 @@
-import { f as p } from "../../index-UdZhlibC.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-UdZhlibC.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-CVYNbZxu.js";
-import { r as v } from "../../Reshape-CEsEp0AI.js";
-import { f as w, p as P } from "../../index-UdZhlibC.js";
-import { e as S } from "../../axis_util-5DTW2tFV.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,4 +1,4 @@
-import { f as l } from "../../index-UdZhlibC.js";
+import { f as l } from "../../index-BoWRt-10.js";
 class u {
   variableNames = ["labels", "logits", "values"];
   outputShape;

package/dist/ops/webgl/gelu.js

@@ -1,5 +1,5 @@
-import { f as a } from "../../index-UdZhlibC.js";
-import { u as s, C as x } from "../../kernel_funcs_utils-CXDy3EN7.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, af as r } from "../../index-UdZhlibC.js";
-import { u as s, l as N } from "../../kernel_funcs_utils-CXDy3EN7.js";
-import { y as l } from "../../shared-DS5waSIY.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-UdZhlibC.js";
-import { r as f } from "../../Reshape-CEsEp0AI.js";
-import { M as U } from "../../mulmat_packed_gpu-IGPBp6h9.js";
-import { m as E } from "../../mat_mul-dwmZz69e.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-UdZhlibC.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-UdZhlibC.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-UdZhlibC.js";
-import { s as x } from "../../sum-NLYbiDag.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-UdZhlibC.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-UdZhlibC.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-UdZhlibC.js";
+import { f as i } from "../../index-BoWRt-10.js";
 class u {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;

package/dist/ops/webgpu/adamAdjust.js

@@ -1,7 +1,7 @@
-import { g as p } from "../../webgpu_program-WaoMq-WD.js";
-import { f as d, c as l } from "../../webgpu_util-DhSeP4b6.js";
-import { f } from "../../index-UdZhlibC.js";
-class c {
+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";
@@ -26,7 +26,7 @@ class c {
                 let m2Hat = moments.y * uniforms.invbeta2;
 
                 let invSqrt = inverseSqrt(max(m2Hat, 1e-30));
-                let invDenom = invSqrt / (1.0 + uniforms.epsilon * invSqrt);
+                let invDenom = invSqrt / fma(uniforms.epsilon, invSqrt, 1.0);
                 let adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);
 
                 setOutputAtIndex(index, adjustedValue);
@@ -35,18 +35,20 @@ class c {
     `;
   }
 }
-function h(t) {
-  const { moments: e, value: a } = t.inputs, { beta1: n, beta2: o, learningRate: i, epsilon: s } = t.attrs, r = t.backend, u = new c(a.shape), m = [
+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 / o] },
-    { type: "float32", data: [i] },
-    { type: "float32", data: [s] }
+    { type: "float32", data: [1 / s] },
+    { type: "float32", data: [o] },
+    { type: "float32", data: [i] }
   ];
   return r.runWebGPUProgram(u, [e, a], "float32", m);
 }
-const v = {
+const b = {
   kernelName: "AdamAdjust",
   backendName: "webgpu",
-  kernelFunc: h
+  kernelFunc: v
 };
-f(v);
+f(b);

package/dist/ops/webgpu/adamMoments.js

@@ -1,7 +1,7 @@
-import { g as m } from "../../webgpu_program-WaoMq-WD.js";
-import { f as u, c as p } from "../../webgpu_util-DhSeP4b6.js";
-import { f as c } from "../../index-UdZhlibC.js";
-class d {
+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";
@@ -25,7 +25,9 @@ class d {
         ${m("index")} {
             if (index < uniforms.size) {
                 let m: vec2<f32> = moments[index];
-                let g: f32 = gradient[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));
@@ -37,15 +39,20 @@ class d {
   }
 }
 function f(t) {
-  const { moments: e, gradient: a } = t.inputs, { beta1: n, beta2: o } = t.attrs, s = t.backend, i = new d(e.shape), r = [
-    { type: "float32", data: [n] },
-    { type: "float32", data: [o] }
+  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 s.runWebGPUProgram(i, [e, a], "float32", r);
+  return r.runWebGPUProgram(s, [e, o], "float32", i);
 }
-const h = {
+const l = {
   kernelName: "AdamMoments",
   backendName: "webgpu",
   kernelFunc: f
 };
-c(h);
+d(l);

package/dist/ops/webgpu/appendCache.js

@@ -1,7 +1,7 @@
-import { g as u } from "../../webgpu_program-WaoMq-WD.js";
-import { f as d, c as m } from "../../webgpu_util-DhSeP4b6.js";
-import { f as l } from "../../index-UdZhlibC.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";
@@ -10,12 +10,12 @@ class f {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "cacheT: i32";
-  constructor(t, a, o, c, i) {
-    const s = Math.min(o + 1, i);
-    this.shaderKey = `AppendCache_${s}`, this.outputShape = [t, a, s, c], this.dispatchLayout = d(this.outputShape), this.dispatch = m(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 `
         ${u("index")} {
             if (index < uniforms.size) {
@@ -26,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);
 
@@ -45,13 +45,16 @@ class f {
         `;
   }
 }
-function T(e) {
-  const { cache: t, item: a } = e.inputs, { maxSize: o, pastLen: c } = e.attrs, i = e.backend, s = t.shape[0], r = t.shape[2], n = t.shape[1], h = new f(s, n, r, a.shape[3], o), p = [{ type: "int32", data: [c] }];
-  return i.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
 };
-l(g);
+l(T);

package/dist/ops/webgpu/attentionMask.js

@@ -1,6 +1,6 @@
-import { f } from "../../index-UdZhlibC.js";
-import { g as m } from "../../webgpu_program-WaoMq-WD.js";
-import { f as k, c as l } from "../../webgpu_util-DhSeP4b6.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;
@@ -14,13 +14,13 @@ class g {
   nh;
   T1;
   T2;
-  constructor(t, e, i, r, s) {
-    if (this.shaderKey = `AttentionMask_${s}`, this.outputShape = [t, e, i, r], this.hs = s, this.nh = e, this.T1 = i, this.T2 = r, this.dispatchLayout = k(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize), s % 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 `
-            ${m("index")} {
+            ${k("index")} {
                 
                 let coords = getCoordsFromIndex(index);
                 let b = coords[0];
@@ -33,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;
@@ -49,17 +50,22 @@ class g {
         `;
   }
 }
-function q(o) {
-  const { q: t, k: e } = o.inputs, { divisor: i, pastLen: r } = o.attrs, s = o.backend, a = t.shape[0], n = t.shape[2], u = e.shape[2], d = t.shape[1], h = t.shape[3], c = new g(a, d, n, u, h), p = [
-    { type: "float32", data: [i] },
-    { type: "int32", data: [r] },
+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 s.runWebGPUProgram(c, [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);

package/dist/ops/webgpu/gatherSub.js

@@ -1,6 +1,6 @@
-import { g as r } from "../../webgpu_program-WaoMq-WD.js";
-import { f as u, c as d } from "../../webgpu_util-DhSeP4b6.js";
-import { f as p } from "../../index-UdZhlibC.js";
+import { g as u } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as h, c as p } from "../../webgpu_util-D____QpY.js";
+import { f as c, a2 as r } from "../../index-BoWRt-10.js";
 class l {
   variableNames = ["labels", "logits", "values"];
   outputShape;
@@ -10,43 +10,30 @@ class l {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = [e], this.dispatchLayout = u(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = [e], this.dispatchLayout = h(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-        ${r("index")} {
+        ${u("index")} {
             if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let idx = i32(getLabelsByOutputIndex(index));
-                let val = getValuesByOutputIndex(index);
-
-                if (isnan(val)) {
-                    setOutputAtIndex(index, 0.0);
-                    return;
-                }
-
-                if (idx < uniforms.logitsShape[1] && idx >= 0) {
-                    let logit = getLogits(coords, idx);
-                    if (isnan(logit)) {
-                        setOutputAtIndex(index, 0.0);
-                        return;
-                    }
-                    setOutputAtIndex(index, val - logit);
-                } else {
-                    setOutputAtIndex(index, 0.0);
-                }
+                let idx = i32(labels[index]);
+                let val = values[index];
+                let logit = logits[index * uniforms.logitsShape[1] + idx];
+                setOutputAtIndex(index, val - logit);
             }
         }
     `;
   }
 }
-function c(t) {
-  const { logits: e, labels: i, values: s } = t.inputs, a = t.backend, n = i.shape[0], o = new l(n);
-  return a.runWebGPUProgram(o, [i, e, s], "float32");
+function d(t) {
+  const { logits: e, labels: a, values: s } = t.inputs, o = t.backend, i = a.shape[0];
+  r(s.shape, [i], "Error in EfficientGatherSub: "), r(a.shape, [i], "Error in EfficientGatherSub: ");
+  const n = new l(i);
+  return o.runWebGPUProgram(n, [a, e, s], "float32");
 }
-const h = {
+const f = {
   kernelName: "EfficientGatherSub",
   backendName: "webgpu",
-  kernelFunc: c
+  kernelFunc: d
 };
-p(h);
+c(f);

package/dist/ops/webgpu/gelu.js

@@ -1,6 +1,6 @@
-import { f as i } from "../../index-UdZhlibC.js";
-import { g as o } from "../../webgpu_program-WaoMq-WD.js";
-import { f as s, c as p } from "../../webgpu_util-DhSeP4b6.js";
+import { f as i } from "../../index-BoWRt-10.js";
+import { g as o } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as s, c as p } from "../../webgpu_util-D____QpY.js";
 const u = 0.7978845608028654, a = 0.044715;
 class c {
   outputShape;

package/dist/ops/webgpu/normRMS.js

@@ -1,5 +1,5 @@
-import { f as p } from "../../webgpu_util-DhSeP4b6.js";
-import { f as s } from "../../index-UdZhlibC.js";
+import { f as n } from "../../webgpu_util-D____QpY.js";
+import { f as p, a2 as h } from "../../index-BoWRt-10.js";
 import { createReduceInfo as u, reduce as c, createReductionShader as m } from "./utils/reductions.js";
 class d {
   outputShape;
@@ -12,7 +12,7 @@ class d {
   inputShape;
   size = !0;
   constructor(e) {
-    this.inputShape = [e.batchSize, e.inSize], this.outputShape = this.inputShape, this.dispatchLayout = p(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
+    this.inputShape = [e.batchSize, e.inSize], this.outputShape = this.inputShape, this.dispatchLayout = n(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
   }
   getUserCode() {
     const e = this.workgroupSize[0];
@@ -29,13 +29,21 @@ class d {
         `);
   }
 }
-function S(t) {
-  const { x: e, gamma: i } = t.inputs, n = t.backend, a = [e, i], o = u(a, -1), r = new d(o);
-  return c(r, a, n);
+function S(a) {
+  const { x: e, gamma: r } = a.inputs, o = a.backend, s = [e, r], t = u(s, -1), i = new d(t);
+  if (h(r.shape, [e.shape[e.shape.length - 1]], "Error in RMSNorm: "), e.shape.length !== 3)
+    throw new Error(`rmsNormGPU: input rank ${e.shape.length} not supported, only rank 3 is supported`);
+  if (t.inSize !== e.shape[2])
+    throw new Error(`rmsNormGPU: reduction size ${t.inSize} does not match expected size ${e.shape[2]}`);
+  if (t.batchSize !== e.shape[0] * e.shape[1])
+    throw new Error(
+      `rmsNormGPU: batch size ${t.batchSize} does not match expected size ${e.shape[0] * e.shape[1]}`
+    );
+  return c(i, s, o);
 }
-const h = {
+const l = {
   kernelName: "RMSNorm",
   backendName: "webgpu",
   kernelFunc: S
 };
-s(h);
+p(l);

package/dist/ops/webgpu/normRMSGrad.js

@@ -1,10 +1,10 @@
-import { f as h, e as S } from "../../index-UdZhlibC.js";
-import { createReduceInfo as g } from "./utils/reductions.js";
-import { f as k } from "../../webgpu_util-DhSeP4b6.js";
-import { g as x } from "../../webgpu_program-WaoMq-WD.js";
-import { s as c } from "../../slice-BHbDHObE.js";
-import { s as z } from "../../sum-NLYbiDag.js";
-class w {
+import { f, a2 as c, e as g } from "../../index-BoWRt-10.js";
+import { createReduceInfo as k } from "./utils/reductions.js";
+import { f as x } from "../../webgpu_util-D____QpY.js";
+import { g as z } from "../../webgpu_program-PFzf1hAQ.js";
+import { s as d } from "../../slice-BiOsknYS.js";
+import { s as w } from "../../sum-66ew2byf.js";
+class y {
   outputShape;
   shaderKey = "RMSNormGrad";
   dispatchLayout;
@@ -16,7 +16,7 @@ class w {
   size = !1;
   rowsPerWorkgroup;
   constructor(r, e = 4) {
-    if (this.shaderKey = `RMSNormGrad_${e}`, this.rowsPerWorkgroup = e, this.inputShape = [r.batchSize, r.inSize], this.outputShape = [r.batchSize + r.batchSize / this.rowsPerWorkgroup, r.inSize], this.dispatchLayout = k(this.outputShape), this.dispatch = [r.batchSize / this.rowsPerWorkgroup, 1, 1], r.batchSize % this.rowsPerWorkgroup !== 0)
+    if (this.shaderKey = `RMSNormGrad_${e}`, this.rowsPerWorkgroup = e, this.inputShape = [r.batchSize, r.inSize], this.outputShape = [r.batchSize + r.batchSize / this.rowsPerWorkgroup, r.inSize], this.dispatchLayout = x(this.outputShape), this.dispatch = [r.batchSize / this.rowsPerWorkgroup, 1, 1], r.batchSize % this.rowsPerWorkgroup !== 0)
       throw new Error(
         `RMSNormGradProgram: batch size ${r.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
       );
@@ -35,7 +35,7 @@ class w {
           var<workgroup> accumulation: array<f32, 1024>;
         `}
 
-          ${x("index")} {
+          ${z("index")} {
             // One workgroup per row (batch).
             let Length = uniforms.reduceSize;
             let BatchSize = uniforms.batchSize;
@@ -105,24 +105,29 @@ class w {
         `;
   }
 }
-function v(t) {
-  const { dy: r, x: e, gamma: o } = t.inputs, i = 4, l = t.backend, a = g([e, o, r], -1), s = new w(a, i), d = [
+function v(o) {
+  const { dy: r, x: e, gamma: t } = o.inputs, i = 4;
+  c(e.shape, r.shape, "Error in RMSNormGrad dy: "), c(t.shape, [e.shape[e.shape.length - 1]], "Error in RMSNormGrad gamma: ");
+  const l = o.backend, a = k([e, t, r], -1), s = new y(a, i), p = [
     { type: "int32", data: [s.inputShape[1]] },
     // Reduce size
     { type: "int32", data: [s.inputShape[0]] }
     // Batch size
-  ], n = S().makeTensorFromTensorInfo(
-    l.runWebGPUProgram(s, [e, o, r], "float32", d)
-  ), u = c(n, [0, 0], [a.batchSize, a.inSize]), m = c(n, [a.batchSize, 0], [a.batchSize / i, a.inSize]);
+  ];
+  if (a.inSize > 1024)
+    throw new Error(`rmsNormGradGPU: inSize ${a.inSize} exceeds max of 1024`);
+  const n = g().makeTensorFromTensorInfo(
+    l.runWebGPUProgram(s, [e, t, r], "float32", p)
+  ), m = d(n, [0, 0], [a.batchSize, a.inSize]), u = d(n, [a.batchSize, 0], [a.batchSize / i, a.inSize]);
   n.dispose();
-  const p = u.reshape(e.shape);
-  u.dispose();
-  const f = z(m, [0]);
-  return m.dispose(), [p, f];
+  const h = m.reshape(e.shape);
+  m.dispose();
+  const S = w(u, [0]);
+  return u.dispose(), [h, S];
 }
-const y = {
+const b = {
   kernelName: "RMSNormGrad",
   backendName: "webgpu",
   kernelFunc: v
 };
-h(y);
+f(b);