@genai-fi/nanogpt 0.7.0 → 0.7.2

package/dist/ops/webgpu/gatherSub.js

@@ -1,6 +1,7 @@
-import { f as u, c as n, g as c } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
-class d {
+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;
   shaderKey = "GatherSub";
@@ -9,29 +10,30 @@ class d {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = [e], this.dispatchLayout = u(this.outputShape), this.dispatch = n(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 `
-        ${c("index")} {
+        ${u("index")} {
             if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let idx = i32(getLabelsByOutputIndex(index));
-                let val = getValuesByOutputIndex(index);
-                let logit = getLogits(coords, idx);
+                let idx = i32(labels[index]);
+                let val = values[index];
+                let logit = logits[index * uniforms.logitsShape[1] + idx];
                 setOutputAtIndex(index, val - logit);
             }
         }
     `;
   }
 }
-function l(t) {
-  const { logits: e, labels: a, values: s } = t.inputs, i = t.backend, o = a.shape[0], r = new d(o);
-  return i.runWebGPUProgram(r, [a, 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: l
+  kernelFunc: d
 };
-p(h);
+c(f);

package/dist/ops/webgpu/gelu.js

@@ -1,5 +1,6 @@
-import { f as i } from "../../index-C0dhsYom.js";
-import { f as o, c as s, g as p } from "../../webgpu_util-g13LvDIv.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;
@@ -10,7 +11,7 @@ class c {
   workgroupSize;
   size = !0;
   constructor(e) {
-    this.workgroupSize = [128, 1, 1], this.outputShape = e, this.dispatchLayout = o(this.outputShape), this.dispatch = s(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_gelu";
+    this.workgroupSize = [128, 1, 1], this.outputShape = e, this.dispatchLayout = s(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_gelu";
   }
   getUserCode() {
     return `
@@ -22,7 +23,7 @@ class c {
         inner = 0.5 * (1.0 + inner);
         return x * inner;
       }
-      ${p("index")} {
+      ${o("index")} {
         if (index < uniforms.size) {
           let a = getAByOutputIndex(index);
           setOutputAtIndex(index, unaryOperation(a));
@@ -51,11 +52,11 @@ class x {
   workgroupSize = [128, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = e, this.dispatchLayout = o(this.outputShape), this.dispatch = s(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = e, this.dispatchLayout = s(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-            ${p("index")} {
+            ${o("index")} {
                 if (index < uniforms.size) {
                     let X  = getXByOutputIndex(index);
                     let x2 = X * X;

package/dist/ops/webgpu/index.js

@@ -1,8 +1,11 @@
 import "./attentionMask.js";
 import "./normRMS.js";
+import "./normRMSGrad.js";
 import "./rope.js";
 import "./appendCache.js";
 import "./scatterSub.js";
 import "./gatherSub.js";
 import "./qkv.js";
 import "./gelu.js";
+import "./adamMoments.js";
+import "./adamAdjust.js";

package/dist/ops/webgpu/normRMS.js

@@ -1,115 +1,49 @@
-import { f as m, c as p, g as c } from "../../webgpu_util-g13LvDIv.js";
-import { f as l, e as k } from "../../index-C0dhsYom.js";
-import { s as M } from "../../sum-B8wEpKsg.js";
-class N {
-  variableNames = ["x", "meanSquare", "gamma"];
+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;
   shaderKey = "RMSNorm";
   dispatchLayout;
   dispatch;
   workgroupSize = [64, 1, 1];
+  variableNames = ["x", "gamma"];
+  uniforms = "reduceSize : i32,";
+  inputShape;
   size = !0;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(e) {
+    this.inputShape = [e.batchSize, e.inSize], this.outputShape = this.inputShape, this.dispatchLayout = n(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
   }
   getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0);
-                let gamma = getGammaByOutputIndex(index);
-                let invRms = inverseSqrt(meanSquare + 1e-8);
-                let normalized = x * invRms;
-                let outVal = normalized * gamma;
-                setOutputAtIndex(index, outVal);
-            }
-        }
-        `;
+    const e = this.workgroupSize[0];
+    return m(e, "mean", `
+            candidate = candidate * candidate;
+        `, `
+            bestValue = inverseSqrt(bestValue + 1e-8);
+        `, `
+            let X = f32(x[offset + k]);
+            let gamma = gamma[k];
+            let normalized = X * bestValue;
+            let outVal = normalized * gamma;
+            result[offset + k] = f32(outVal);
+        `);
   }
 }
-function b(s) {
-  const { x: t, gamma: e } = s.inputs, a = s.backend, o = t.shape[0], n = t.shape[1], i = t.shape[2], u = t.square().mean(-1, !0), r = new N(o, n, i);
-  return a.runWebGPUProgram(r, [t, u, e], "float32");
+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 z = {
+const l = {
   kernelName: "RMSNorm",
   backendName: "webgpu",
-  kernelFunc: b
+  kernelFunc: S
 };
-l(z);
-class R {
-  variableNames = ["x", "meanSquare", "dyGamma", "dyXMean"];
-  outputShape;
-  shaderKey = "RMSNormGradX";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  C;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize), this.C = a;
-  }
-  getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0) + 1e-8;
-                let dyGamma = getDyGammaByOutputIndex(index);
-                let dyXMean = getDyXMean(coords[0], coords[1], 0) / ${this.C}.0;
-                let invRms = inverseSqrt(meanSquare);
-                let dx = dyGamma * invRms - x * dyXMean * invRms / meanSquare;
-                setOutputAtIndex(index, dx);
-            }
-        }
-        `;
-  }
-}
-class v {
-  variableNames = ["x", "meanSquare", "dy"];
-  outputShape;
-  shaderKey = "RMSNormGradGamma";
-  dispatchLayout;
-  dispatch;
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  constructor(t, e, a) {
-    this.outputShape = [t, e, a], this.dispatchLayout = m(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
-  }
-  getUserCode() {
-    return `
-        ${c("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index);
-                let x = getXByOutputIndex(index);
-                let meanSquare = getMeanSquare(coords[0], coords[1], 0) + 1e-8;
-                let dy = getDyByOutputIndex(index);
-                let invRms = inverseSqrt(meanSquare);
-                let dGamma = dy * (x * invRms);
-                setOutputAtIndex(index,dGamma);
-            }
-        }
-        `;
-  }
-}
-function I(s) {
-  const { dy: t, x: e, gamma: a } = s.inputs, o = s.backend, n = e.shape[0], i = e.shape[1], u = e.shape[2], r = t.mul(a), h = r.mul(e), g = h.sum(-1, !0);
-  h.dispose();
-  const S = e.square(), d = S.mean(-1, !0);
-  S.dispose();
-  const y = new R(n, i, u), G = o.runWebGPUProgram(y, [e, d, r, g], "float32");
-  r.dispose(), g.dispose();
-  const q = new v(n, i, u), x = o.runWebGPUProgram(q, [e, d, t], "float32");
-  d.dispose();
-  const f = M(k().makeTensorFromTensorInfo(x), [0, 1]);
-  return o.disposeData(x), [G, f];
-}
-const P = {
-  kernelName: "RMSNormGrad",
-  backendName: "webgpu",
-  kernelFunc: I
-};
-l(P);
+p(l);

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

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

package/dist/ops/webgpu/normRMSGrad.js

@@ -0,0 +1,133 @@
+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;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  variableNames = ["x", "gamma", "dy"];
+  uniforms = "reduceSize : i32, batchSize: i32";
+  inputShape;
+  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 = 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}`
+      );
+    if (r.inSize > 1024)
+      throw new Error(`RMSNormGradProgram: inSize ${r.inSize} exceeds max of 1024`);
+  }
+  getUserCode() {
+    const r = this.workgroupSize[0], e = this.rowsPerWorkgroup;
+    return `
+          fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+          }
+
+          ${`
+          var<workgroup> partials : array<vec2<f32>, ${r}>;
+          var<workgroup> accumulation: array<f32, 1024>;
+        `}
+
+          ${z("index")} {
+            // One workgroup per row (batch).
+            let Length = uniforms.reduceSize;
+            let BatchSize = uniforms.batchSize;
+            for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                accumulation[k] = 0.0;
+            }
+
+            for (var rowOff = 0; rowOff < ${e}; rowOff = rowOff + 1) {
+                let row = i32(workgroupId.x) * ${e} + rowOff;
+                let offset = row * Length;
+                
+                var sum_x2 = 0.0;
+                var sum_dygx = 0.0;
+
+                for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                    let X = f32(x[offset + k]);
+                    let DY = f32(dy[offset + k]);
+                    let G  = f32(gamma[k]);
+                    sum_x2 = fma(X, X, sum_x2);
+                    sum_dygx = fma(DY * G, X, sum_dygx);
+                }
+
+                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
+                workgroupBarrier();
+
+                var reduceSize = min(u32(Length), ${r}u);
+                for (var currentSize = reduceSize / 2u; reduceSize > 1u; currentSize = reduceSize / 2u) {
+                    let interval = DIV_CEIL(reduceSize, 2u);
+                    if (localId.x < currentSize) {
+                        partials[localId.x] = partials[localId.x] + partials[localId.x + interval];
+                    }
+                    reduceSize = interval;
+                    workgroupBarrier();
+                }
+
+                let invN = 1.0 / f32(Length);
+                let mean_x2   = fma(partials[0].x, invN, 1e-8);
+                let mean_dygx = partials[0].y * invN;
+
+                let invRMS = inverseSqrt(mean_x2);
+                let scale = (mean_dygx / (mean_x2)) * invRMS;
+
+                // write dx and dGamma.
+                for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                    let X  = f32(x[offset + k]);
+                    let DY = f32(dy[offset + k]);
+                    let G  = f32(gamma[k]);
+
+                    let dyGamma = DY * G;
+                    let dx = fma(dyGamma, invRMS, -X * scale);
+
+                    result[offset + k] = dx;
+
+                    // dGamma
+                    accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
+                }
+
+                workgroupBarrier();
+            }
+
+            // Write out the partially accumulated dGamma
+            let outDgBase = BatchSize * Length + i32(workgroupId.x) * Length;
+            for (var k = i32(localId.x); k < Length; k = k + ${r}) {
+                result[outDgBase + k] = accumulation[k];
+            }
+          }
+        `;
+  }
+}
+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
+  ];
+  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 h = m.reshape(e.shape);
+  m.dispose();
+  const S = w(u, [0]);
+  return u.dispose(), [h, S];
+}
+const b = {
+  kernelName: "RMSNormGrad",
+  backendName: "webgpu",
+  kernelFunc: v
+};
+f(b);

package/dist/ops/webgpu/qkv.js

@@ -1,5 +1,6 @@
-import { f as c, c as d, g as h } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
+import { g as h } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as c, c as d } from "../../webgpu_util-D____QpY.js";
+import { f as p, a2 as m } from "../../index-BoWRt-10.js";
 class l {
   variableNames = ["x", "kernel"];
   outputShape;
@@ -10,8 +11,8 @@ class l {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(t, e, o, s) {
-    const r = s / e;
-    this.outputShape = [t, e, o, r], this.dispatchLayout = c(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    const n = s / e;
+    this.shaderKey = `QKV_${e}_${n}`, this.outputShape = [t, e, o, n], this.dispatchLayout = c(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     const t = this.outputShape[1], e = this.outputShape[3], o = t * e;
@@ -25,13 +26,14 @@ class l {
                 let d = coords[3];
 
                 // Compute output channel index in fused kernel
-                let out_offset = uniforms.mode * ${t} * ${e} + h * ${e} + d;
+                let out_offset = uniforms.mode * ${o} + h * ${e} + d;
 
                 var sum = 0.0;
+                let baseX = b * uniforms.xShape[1] * uniforms.xShape[2] + t * uniforms.xShape[2];
                 for (var c = 0; c < ${o}; c += 1) {
-                    let xval = getX(b, t, c); // fetch from x
-                    let kval = getKernel(c, out_offset); // fetch from kernel
-                    sum += xval * kval;
+                    let xval = x[baseX + c];
+                    let kval = getKernel(c, out_offset);
+                    sum = fma(xval, kval, sum);
                 }
 
                 setOutputAtIndex(index, sum);
@@ -40,17 +42,20 @@ class l {
         `;
   }
 }
-function m(a) {
-  const { x: t, kernel: e } = a.inputs, { heads: o } = a.attrs, s = a.backend, r = t.shape[0], i = t.shape[1], u = t.shape[2], n = new l(r, o, i, u);
+function f(r) {
+  const { x: t, kernel: e } = r.inputs, { heads: o } = r.attrs, s = r.backend, n = t.shape[0], u = t.shape[1], a = t.shape[2];
+  if (m(e.shape, [a, 3 * a], "Error in QKV: "), a % o !== 0)
+    throw new Error(`Channel dimension ${a} must be divisible by number of heads ${o} in QKV.`);
+  const i = new l(n, o, u, a);
   return [
-    s.runWebGPUProgram(n, [t, e], "float32", [{ type: "int32", data: [0] }]),
-    s.runWebGPUProgram(n, [t, e], "float32", [{ type: "int32", data: [1] }]),
-    s.runWebGPUProgram(n, [t, e], "float32", [{ type: "int32", data: [2] }])
+    s.runWebGPUProgram(i, [t, e], "float32", [{ type: "int32", data: [0] }]),
+    s.runWebGPUProgram(i, [t, e], "float32", [{ type: "int32", data: [1] }]),
+    s.runWebGPUProgram(i, [t, e], "float32", [{ type: "int32", data: [2] }])
   ];
 }
-const f = {
+const b = {
   kernelName: "QKV",
   backendName: "webgpu",
-  kernelFunc: m
+  kernelFunc: f
 };
-p(f);
+p(b);

package/dist/ops/webgpu/rope.js

@@ -1,6 +1,7 @@
-import { f as h, c as l, g as m } from "../../webgpu_util-g13LvDIv.js";
-import { f as g } from "../../index-C0dhsYom.js";
-class f {
+import { g as c } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as m, c as l } from "../../webgpu_util-D____QpY.js";
+import { f as x, a2 as f } from "../../index-BoWRt-10.js";
+class S {
   variableNames = ["x", "sin", "cos"];
   outputShape;
   shaderKey = "Rope";
@@ -9,13 +10,13 @@ class f {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "pastLen: i32";
-  constructor(t, o, s, n) {
-    this.outputShape = [t, o, s, n], this.dispatchLayout = h(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(t, e, a, s) {
+    this.shaderKey = `Rope_${s}`, this.outputShape = [t, e, a, s], this.dispatchLayout = m(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     const t = this.outputShape[3];
     return `
-        ${m("index")} {
+        ${c("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -27,27 +28,31 @@ class f {
 
                 var outVal = 0.0;
 
+                let xIdx = b * uniforms.outShapeStrides[0] +
+                    h * uniforms.outShapeStrides[1] +
+                    t * uniforms.outShapeStrides[2] +
+                    d;
+
                 if (d < rotaryDim) {
-                    let pairIdx = d / 2;
-                    let cos = getCos(t + uniforms.pastLen, pairIdx, 0);
-                    let sin = getSin(t + uniforms.pastLen, pairIdx, 0);
+                    let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d / 2;
+                    let cos = cos[idx];
+                    let sin = sin[idx];
 
-                    let ownX = getX(b, h, t, d) * cos;
+                    let ownX = x[xIdx] * cos;
+                    var evenOdd = 0.0;
 
                     if (d % 2 == 0) {
                         // even index
-                        let even = ownX;
-                        let odd = getX(b, h, t, d + 1);
-                        outVal = even - odd * sin;
+                        evenOdd = -x[xIdx + 1];
                     } else {
                         // odd index
-                        let even = getX(b, h, t, d - 1);
-                        let odd = ownX;
-                        outVal = even * sin + odd;
+                        evenOdd = x[xIdx - 1];
                     }
+
+                    outVal = fma(evenOdd, sin, ownX);
                 } else {
                     // pass through for non-rotary dims
-                    outVal = getX(b, h, t, d);
+                    outVal = x[xIdx];
                 }
 
                 setOutputAtIndex(index, outVal);
@@ -56,13 +61,22 @@ class f {
         `;
   }
 }
-function b(e) {
-  const { x: t, sin: o, cos: s } = e.inputs, { pastLen: n } = e.attrs, a = e.backend, i = t.shape[0], r = t.shape[1], d = t.shape[2], p = t.shape[3], u = new f(i, r, d, p), c = [{ type: "int32", data: [n] }];
-  return a.runWebGPUProgram(u, [t, o, s], "float32", c);
+function g(o) {
+  const { x: t, sin: e, cos: a } = o.inputs, { pastLen: s } = o.attrs, i = o.backend, p = t.shape[0], d = t.shape[1], n = t.shape[2], r = t.shape[3];
+  if (f(e.shape, a.shape, "Error in Rope: "), e.shape[0] < n + s)
+    throw new Error(
+      `Sin tensor shape ${e.shape} is not compatible with seqLength ${n} and pastLen ${s}.`
+    );
+  if (e.shape[1] * 2 < r)
+    throw new Error(`Sin tensor shape ${e.shape} is not compatible with feature dimension ${r}.`);
+  if (e.shape.length !== 3)
+    throw new Error(`Sin tensor must be 3-dimensional, but got shape ${e.shape}.`);
+  const h = new S(p, d, n, r), u = [{ type: "int32", data: [s] }];
+  return i.runWebGPUProgram(h, [t, e, a], "float32", u);
 }
-const x = {
+const b = {
   kernelName: "Rope",
   backendName: "webgpu",
-  kernelFunc: b
+  kernelFunc: g
 };
-g(x);
+x(b);

package/dist/ops/webgpu/scatterSub.js

@@ -1,6 +1,7 @@
-import { f as c, c as u, g as d } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
-class h {
+import { g as p } from "../../webgpu_program-PFzf1hAQ.js";
+import { f as u, c as d } from "../../webgpu_util-D____QpY.js";
+import { f as h, a2 as o } from "../../index-BoWRt-10.js";
+class b {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;
   shaderKey = "ScatterSub";
@@ -9,29 +10,31 @@ class h {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(t, e) {
-    this.outputShape = [t, e], this.dispatchLayout = c(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = [t, e], this.dispatchLayout = u(this.outputShape), this.dispatch = d(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-            ${d("index")} {
+            ${p("index")} {
                 if (index < uniforms.size) {
                     let coords = getCoordsFromIndex(index); // [batch, depth]
-                    let idx = i32(getLabels(coords[0]));
-                    let prob = getSoftmaxProbsByOutputIndex(index);
-                    let dy = getDy(coords[0]);
+                    let idx = i32(labels[coords[0]]);
+                    let prob = softmaxProbs[index];
+                    let dy = dy[coords[0]];
                     setOutputAtIndex(index, select(prob, prob - 1.0, idx == coords[1]) * dy);
                 }
             }
             `;
   }
 }
-function b(o) {
-  const { logits: t, labels: e, dy: s } = o.inputs, a = o.backend, r = e.shape[0], i = t.shape[1], n = new h(r, i);
-  return a.runWebGPUProgram(n, [e, t, s], "float32");
+function f(a) {
+  const { logits: t, labels: e, dy: s } = a.inputs, c = a.backend, r = e.shape[0], i = t.shape[1];
+  o(s.shape, [r], "Error in EfficientScatterSub dy: "), o(t.shape, [r, i], "Error in EfficientScatterSub logits: "), o(e.shape, [r], "Error in EfficientScatterSub labels: ");
+  const n = new b(r, i);
+  return c.runWebGPUProgram(n, [e, t, s], "float32");
 }
 const l = {
   kernelName: "EfficientScatterSub",
   backendName: "webgpu",
-  kernelFunc: b
+  kernelFunc: f
 };
-p(l);
+h(l);

package/dist/ops/webgpu/utils/reductions.d.ts

@@ -0,0 +1,9 @@
+import { backend_util, TensorInfo } from '@tensorflow/tfjs-core';
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
+export interface ReduceWebGPUProgram extends WebGPUProgram {
+    inputShape: number[];
+}
+export declare function createReductionShader(workgroupSizeX: number, reductionOp: 'mean' | 'sum', inputSnippet: string, reducedSnippet: string, outputSnippet: string): string;
+export declare function createReduceInfo(inputs: TensorInfo[], axis: number | number[]): backend_util.ReduceInfo;
+export declare function reduce(program: ReduceWebGPUProgram, inputs: TensorInfo[], backend: WebGPUBackend): TensorInfo;