@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/ops/webgpu/normRMSGrad.js

@@ -1,10 +1,15 @@
-import { f, a4 as c, e as g } from "../../index-BzFyqcy-.js";
-import { createReduceInfo as k } from "./utils/reductions.js";
-import { f as x } from "../../webgpu_util-pLEV9tks.js";
-import { g as z } from "../../webgpu_program-DkQJOJSd.js";
-import { s as d } from "../../slice-d0Vo9XTN.js";
-import { s as w } from "../../sum-CJ0ULhmt.js";
-class y {
+import { e as _ } from "../../index-ZyQhjEPo.js";
+import { createReduceInfo as D } from "./utils/reductions.js";
+import { f as X } from "../../webgpu_util-BBCnKm2X.js";
+import { e as $ } from "../../webgpu_program-Cigz-7RF.js";
+import { j as z } from "../../tensor-DdQUJZlz.js";
+import { p as k, u as M } from "../../pack16-CFUqumar.js";
+import { isPackedTensor as h } from "../../utilities/packed.js";
+import { reshape16 as R } from "../reshape16.js";
+import { sum16 as L } from "../sum16.js";
+import { slice16 as w } from "../slice16.js";
+import { r as P } from "../../tensor_util-DV-FP5Q3.js";
+class N {
   outputShape;
   shaderKey = "RMSNormGrad";
   dispatchLayout;
@@ -15,53 +20,96 @@ class y {
   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)
+  packed = !1;
+  outputComponent;
+  constructor(a, e = 4, o = !1) {
+    if (this.packed = o, this.shaderKey = `RMSNormGrad_${e}`, this.rowsPerWorkgroup = e, this.inputShape = [a.batchSize, a.inSize], this.outputShape = [a.batchSize + a.batchSize / this.rowsPerWorkgroup, a.inSize], this.dispatchLayout = X(this.outputShape), this.dispatch = [a.batchSize / this.rowsPerWorkgroup, 1, 1], a.batchSize % this.rowsPerWorkgroup !== 0)
       throw new Error(
-        `RMSNormGradProgram: batch size ${r.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
+        `RMSNormGradProgram: batch size ${a.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
       );
-    if (r.inSize > 1024)
-      throw new Error(`RMSNormGradProgram: inSize ${r.inSize} exceeds max of 1024`);
+    if (a.inSize > 1024)
+      throw new Error(`RMSNormGradProgram: inSize ${a.inSize} exceeds max of 1024`);
   }
   getUserCode() {
-    const r = this.workgroupSize[0], e = this.rowsPerWorkgroup;
+    const a = this.workgroupSize[0], e = this.rowsPerWorkgroup, o = `
+          var<workgroup> partials : array<vec2<f32>, ${a}>;
+          var<workgroup> accumulation: array<${this.packed ? "vec2<f32>" : "f32"}, 1024>;
+        `, n = this.packed ? `
+                let X = unpack2x16float(u32(x[offset + k]));
+                let DY = unpack2x16float(u32(dy[offset + k]));
+                let G  = unpack2x16float(u32(gamma[k]));
+                sum_x2 = fma(X.x, X.x, sum_x2);
+                sum_x2 = fma(X.y, X.y, sum_x2);
+                sum_dygx = fma(DY.x * G.x, X.x, sum_dygx);
+                sum_dygx = fma(DY.y * G.y, X.y, sum_dygx);
+        ` : `
+                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);
+            `, s = this.packed ? `
+                let X  = unpack2x16float(u32(x[offset + k]));
+                let DY = unpack2x16float(u32(dy[offset + k]));
+                let G  = unpack2x16float(u32(gamma[k]));
+
+                let dyGamma = DY * G;
+                let dx = vec2<f32>(
+                    fma(dyGamma.x, invRMS, -X.x * scale),
+                    fma(dyGamma.y, invRMS, -X.y * scale)
+                );
+
+                result[offset + k] = i32(pack2x16float(dx));
+
+                // dGamma
+                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
+        ` : `
+                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]);
+        `, i = this.packed ? `
+                result[outDgBase + k] = i32(pack2x16float(accumulation[k]));
+        ` : `
+                result[outDgBase + k] = accumulation[k];
+        `;
     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>;
-        `}
+          ${o}
 
-          ${z("index")} {
+          ${$("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 k = i32(localId.x); k < Length; k = k + ${a}) {
+                accumulation[k] = ${this.packed ? "vec2<f32>(0.0f)" : "0.0f"};
             }
 
             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);
+                var sum_x2 = 0.0f;
+                var sum_dygx = 0.0f;
+
+                for (var k = i32(localId.x); k < Length; k = k + ${a}) {
+                    ${n}
                 }
 
                 partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
                 workgroupBarrier();
 
-                var reduceSize = min(u32(Length), ${r}u);
+                var reduceSize = min(u32(Length), ${a}u);
                 for (var currentSize = reduceSize / 2u; reduceSize > 1u; currentSize = reduceSize / 2u) {
                     let interval = DIV_CEIL(reduceSize, 2u);
                     if (localId.x < currentSize) {
@@ -71,7 +119,7 @@ class y {
                     workgroupBarrier();
                 }
 
-                let invN = 1.0 / f32(Length);
+                let invN = 1.0f / f32(${this.packed ? "Length * 2" : "Length"});
                 let mean_x2   = fma(partials[0].x, invN, 1e-8);
                 let mean_dygx = partials[0].y * invN;
 
@@ -79,18 +127,8 @@ class y {
                 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]);
+                for (var k = i32(localId.x); k < Length; k = k + ${a}) {
+                    ${s}
                 }
 
                 workgroupBarrier();
@@ -98,36 +136,47 @@ class y {
 
             // 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];
+            for (var k = i32(localId.x); k < Length; k = k + ${a}) {
+                ${i}
             }
           }
         `;
   }
 }
-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]] },
+function W(p) {
+  const { dy: a, x: e, gamma: o } = p.inputs, n = 4;
+  z(e.shape, a.shape, "Error in RMSNormGrad dy: ");
+  const s = h(e), i = h(o), u = h(a), r = s || i || u, m = !r || s ? e : k(e), c = !r || i ? o : k(o), d = !r || u ? a : k(a);
+  z(c.shape, [m.shape[m.shape.length - 1]], "Error in RMSNormGrad gamma: ");
+  const G = p.backend, t = D([m, c, d], -1), f = new N(t, n, r), v = [
+    { type: "int32", data: [f.inputShape[1]] },
     // Reduce size
-    { type: "int32", data: [s.inputShape[0]] }
+    { type: "int32", data: [f.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];
+  if (t.inSize > 1024)
+    throw new Error(`rmsNormGradGPU: inSize ${t.inSize} exceeds max of 1024`);
+  const x = G.runWebGPUProgram(
+    f,
+    [m, c, d],
+    r ? "int32" : "float32",
+    v
+  );
+  x.packed = r, r && !s && m.dispose(), r && !i && c.dispose(), r && !u && d.dispose();
+  const l = _().makeTensorFromTensorInfo(x), S = w(l, [0, 0], [t.batchSize, t.inSize]), g = w(
+    l,
+    [t.batchSize, 0],
+    [t.batchSize / n, t.inSize]
+  );
+  l.dispose();
+  const b = R(S, e.shape);
+  S.dispose();
+  const y = L(g, [0]);
+  return g.dispose(), [b, !r || i ? y : M(y)];
 }
-const b = {
+const Y = {
   kernelName: "RMSNormGrad",
   backendName: "webgpu",
-  kernelFunc: v
+  kernelFunc: W
 };
-f(b);
+P(Y);

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

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

package/dist/ops/webgpu/pack16.js

@@ -0,0 +1,19 @@
+import "../../index-ZyQhjEPo.js";
+import c from "./pack16_program.js";
+import { r as p } from "../../tensor_util-DV-FP5Q3.js";
+function m(n) {
+  const { x: e } = n.inputs, { scaling: t, padding: r } = n.attrs, i = n.backend;
+  if (e.shape[e.shape.length - 1] % 2 !== 0)
+    throw new Error("Last dimension of input tensor must be even to use Pack16.");
+  n.attrs && (n.attrs.originalShape = e.shape);
+  const a = new c(e.shape, r), o = t !== 1;
+  o && a.useScaling();
+  const s = [{ type: "float32", data: [t] }];
+  return i.runWebGPUProgram(a, [e], "int32", o ? s : void 0);
+}
+const u = {
+  kernelName: "Pack16",
+  backendName: "webgpu",
+  kernelFunc: m
+};
+p(u);

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

@@ -0,0 +1,19 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class PackProgram implements WebGPUProgram {
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    variableNames: string[];
+    uniforms?: string;
+    size: boolean;
+    outputComponent: number;
+    scaling: boolean;
+    padding: number;
+    constructor(outShape: number[], padding?: number);
+    useScaling(): void;
+    getUserCode(): string;
+}

package/dist/ops/webgpu/pack16_program.js

@@ -0,0 +1,92 @@
+import { f as o, c as a } from "../../webgpu_util-BBCnKm2X.js";
+import { e as s } from "../../webgpu_program-Cigz-7RF.js";
+class h {
+  outputShape;
+  shaderKey = "Pack16";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  variableNames = ["x"];
+  uniforms;
+  size = !0;
+  outputComponent = 4;
+  scaling = !1;
+  padding = 0;
+  constructor(t, e = 0) {
+    if (t[t.length - 1] % 2 !== 0 && e === 0)
+      throw new Error("Last dimension of output shape must be even to use Pack16.");
+    if (e % 4 !== 0)
+      throw new Error("Padding must be a multiple of 4 to use Pack16.");
+    if (this.outputShape = [...t.slice(0, -1), t[t.length - 1]], e > 0) {
+      this.shaderKey += `_Padded${e}`, this.padding = e;
+      for (let i = this.outputShape.length - 2; i < this.outputShape.length; i++)
+        this.outputShape[i] % this.padding !== 0 && (this.outputShape[i] += this.padding - this.outputShape[i] % this.padding);
+      this.outputComponent = 1;
+    }
+    this.outputShape[this.outputShape.length - 1] /= 2, this.outputShape[this.outputShape.length - 1] % this.outputComponent !== 0 && (this.outputComponent = 1), this.dispatchLayout = o(this.outputShape), this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+      this.outputComponent,
+      1,
+      1
+    ]);
+  }
+  useScaling() {
+    this.shaderKey += "_Scaled", this.uniforms = "scaling : f32,", this.scaling = !0;
+  }
+  getUserCode() {
+    if (this.padding > 0 && this.outputComponent === 1) {
+      const t = this.outputShape.length;
+      return `
+                ${s("index")} {
+                    if (index < uniforms.size) {
+                        var coords = getCoordsFromIndex(index);
+                        coords[${t} - 1] = coords[${t} - 1] * 2;
+                        let row = coords[${t} - 2];
+                        let col = coords[${t} - 1];
+                        let width = uniforms.xShape[${t} - 1];
+                        let height = uniforms.xShape[${t} - 2];
+                    
+                        var value1 = 0.0f;
+                        if (col < width && row < height) {
+                            let baseInputIndex = getIndexFromCoords${t}D(coords, uniforms.xShape);
+                            value1 = x[baseInputIndex] ${this.scaling ? "* uniforms.scaling" : ""};
+                        }
+                        var value2 = 0.0f;
+                        if (col + 1 < width && row < height) {
+                            coords[${t} - 1] = coords[${t} - 1] + 1;
+                            let baseInputIndex = getIndexFromCoords${t}D(coords, uniforms.xShape);
+                            value2 = x[baseInputIndex] ${this.scaling ? "* uniforms.scaling" : ""};
+                        }
+                        let packed = i32(pack2x16float(vec2<f32>(value1, value2)));
+                        result[index] = packed;
+                    }
+                }`;
+    }
+    return this.outputComponent === 1 ? `
+                ${s("index")} {
+                    if (index < uniforms.size) {
+                        let baseInputIndex =  index * 2;
+                        let x1 = x[baseInputIndex] ${this.scaling ? "* uniforms.scaling" : ""};
+                        let x2 = x[baseInputIndex + 1] ${this.scaling ? "* uniforms.scaling" : ""};
+                        let packed = i32(pack2x16float(vec2<f32>(x1, x2)));
+                        result[index] = packed;
+                    }
+                }` : `
+                ${s("index")} {
+                    if (index < uniforms.size) {
+                        let baseInputIndex =  index * 2;
+                        let x1 = x[baseInputIndex] ${this.scaling ? "* uniforms.scaling" : ""};
+                        let x2 = x[baseInputIndex + 1] ${this.scaling ? "* uniforms.scaling" : ""};
+                        let packed = vec4<i32>(
+                            i32(pack2x16float(vec2<f32>(x1.x, x1.y))),
+                            i32(pack2x16float(vec2<f32>(x1.z, x1.w))),
+                            i32(pack2x16float(vec2<f32>(x2.x, x2.y))),
+                            i32(pack2x16float(vec2<f32>(x2.z, x2.w)))
+                        );
+                        result[index] = packed;
+                    }
+                }`;
+  }
+}
+export {
+  h as default
+};

package/dist/ops/webgpu/qkv.js

@@ -1,61 +1,26 @@
-import { g as h } from "../../webgpu_program-DkQJOJSd.js";
-import { f as c, c as d } from "../../webgpu_util-pLEV9tks.js";
-import { f as p, a4 as m } from "../../index-BzFyqcy-.js";
-class l {
-  variableNames = ["x", "kernel"];
-  outputShape;
-  shaderKey = "QKV";
-  dispatchLayout;
-  dispatch;
-  uniforms = "mode: i32";
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  constructor(t, e, o, s) {
-    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;
-    return `
-        ${h("index")} {
-            if (index < uniforms.size) {
-                let coords = getCoordsFromIndex(index); // [b, h, t, d]
-                let b = coords[0];
-                let h = coords[1];
-                let t = coords[2];
-                let d = coords[3];
-
-                // Compute output channel index in fused kernel
-                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 = x[baseX + c];
-                    let kval = getKernel(c, out_offset);
-                    sum = fma(xval, kval, sum);
-                }
-
-                setOutputAtIndex(index, sum);
-            }
-        }
-        `;
-  }
-}
-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(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] }])
+import "../../index-ZyQhjEPo.js";
+import { j as h } from "../../tensor-DdQUJZlz.js";
+import { b as f } from "../../matMul16--R5hOwDG.js";
+import { slice16 as a } from "../slice16.js";
+import { isPackedTensor as l } from "../../utilities/packed.js";
+import { r as u } from "../../tensor_util-DV-FP5Q3.js";
+function k(i) {
+  const { x: r, kernel: c } = i.inputs, { heads: e } = i.attrs, t = r.shape[0], n = r.shape[1], s = r.shape[2], m = l(r);
+  if (h(c.shape, [m ? s * 2 : s, 3 * s], "Error in QKV: "), s % e !== 0)
+    throw new Error(`Channel dimension ${s} must be divisible by number of heads ${e} in QKV.`);
+  const o = f(r, c, !1, !1, {
+    forceOutputShape: [t, n, 3 * e, s / e],
+    perm: [0, 2, 1, 3]
+  }), p = [
+    a(o, [0, 0, 0, 0], [t, e, n, s / e]),
+    a(o, [0, e, 0, 0], [t, e, n, s / e]),
+    a(o, [0, 2 * e, 0, 0], [t, e, n, s / e])
   ];
+  return o.dispose(), p;
 }
 const b = {
   kernelName: "QKV",
   backendName: "webgpu",
-  kernelFunc: f
+  kernelFunc: k
 };
-p(b);
+u(b);

package/dist/ops/webgpu/rope.js

@@ -1,7 +1,10 @@
-import { g as c } from "../../webgpu_program-DkQJOJSd.js";
-import { f as m, c as l } from "../../webgpu_util-pLEV9tks.js";
-import { f as x, a4 as f } from "../../index-BzFyqcy-.js";
-class S {
+import { isPackedTensor as w } from "../../utilities/packed.js";
+import { e as x } from "../../webgpu_program-Cigz-7RF.js";
+import { f as l, c as m } from "../../webgpu_util-BBCnKm2X.js";
+import "../../index-ZyQhjEPo.js";
+import { j as b } from "../../tensor-DdQUJZlz.js";
+import { r as v } from "../../tensor_util-DV-FP5Q3.js";
+class k {
   variableNames = ["x", "sin", "cos"];
   outputShape;
   shaderKey = "Rope";
@@ -10,13 +13,13 @@ class S {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "pastLen: i32";
-  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);
+  constructor(e, s, r, t) {
+    this.shaderKey = `Rope_${t}`, this.outputShape = [e, s, r, t], this.dispatchLayout = l(this.outputShape), this.dispatch = m(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
-    const t = this.outputShape[3];
+    const e = this.outputShape[3];
     return `
-        ${c("index")} {
+        ${x("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -24,7 +27,7 @@ class S {
                 let t = coords[2];
                 let d = coords[3];
 
-                let rotaryDim = ${t};
+                let rotaryDim = ${e};
 
                 var outVal = 0.0;
 
@@ -61,22 +64,74 @@ class S {
         `;
   }
 }
-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)
+class L {
+  variableNames = ["x", "sin", "cos"];
+  outputShape;
+  shaderKey = "Rope";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "pastLen: i32";
+  constructor(e, s, r, t) {
+    this.shaderKey = `Rope_${t}`, this.outputShape = [e, s, r, t / 2], this.dispatchLayout = l(this.outputShape), this.dispatch = m(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+        ${x("index")} {
+            if (index < uniforms.size) {
+                let coords = getCoordsFromIndex(index); // [b, h, t, d]
+                let b = coords[0];
+                let h = coords[1];
+                let t = coords[2];
+                let d = coords[3];
+
+                var outVal = vec2<f32>(0.0, 0.0);
+
+                let xIdx = b * uniforms.outShapeStrides[0] +
+                    h * uniforms.outShapeStrides[1] +
+                    t * uniforms.outShapeStrides[2] +
+                    d;
+
+                let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d;
+                let cos = cos[idx];
+                let sin = sin[idx];
+
+                let xPair = unpack2x16float(u32(x[xIdx]));
+                let ownX = vec2<f32>(xPair.x * cos, xPair.y * cos);
+
+                let evenOdd = vec2<f32>(
+                    -xPair.y,
+                    xPair.x
+                );
+
+                outVal = vec2<f32>(
+                    fma(evenOdd.x, sin, ownX.x),
+                    fma(evenOdd.y, sin, ownX.y)
+                );
+
+                result[index] = i32(pack2x16float(outVal));
+            }
+        }
+        `;
+  }
+}
+function P(i) {
+  const { x: e } = i.inputs, { pastLen: s, negSin: r, ropeCache: t } = i.attrs, f = i.backend, a = w(e), p = e.shape[0], h = e.shape[1], n = e.shape[2], d = a ? e.shape[3] * 2 : e.shape[3], o = r ? t.getNegSin() : t.getSin(), u = t.getCos();
+  if (b(o.shape, u.shape, "Error in Rope: "), o.shape[0] < n + s)
     throw new Error(
-      `Sin tensor shape ${e.shape} is not compatible with seqLength ${n} and pastLen ${s}.`
+      `Sin tensor shape ${o.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);
+  if (o.shape[1] * 2 < d)
+    throw new Error(`Sin tensor shape ${o.shape} is not compatible with feature dimension ${d}.`);
+  if (o.shape.length !== 3)
+    throw new Error(`Sin tensor must be 3-dimensional, but got shape ${o.shape}.`);
+  const S = a ? new L(p, h, n, d) : new k(p, h, n, d), g = [{ type: "int32", data: [s] }], y = a ? "int32" : e.dtype, c = f.runWebGPUProgram(S, [e, o, u], y, g);
+  return c.packed = a, c;
 }
-const b = {
+const $ = {
   kernelName: "Rope",
   backendName: "webgpu",
-  kernelFunc: g
+  kernelFunc: P
 };
-x(b);
+v($);

package/dist/ops/webgpu/scatterSub.js

@@ -1,6 +1,8 @@
-import { g as p } from "../../webgpu_program-DkQJOJSd.js";
-import { f as u, c as d } from "../../webgpu_util-pLEV9tks.js";
-import { f as h, a4 as o } from "../../index-BzFyqcy-.js";
+import { e as p } from "../../webgpu_program-Cigz-7RF.js";
+import { f as u, c as d } from "../../webgpu_util-BBCnKm2X.js";
+import "../../index-ZyQhjEPo.js";
+import { j as s } from "../../tensor-DdQUJZlz.js";
+import { r as h } from "../../tensor_util-DV-FP5Q3.js";
 class b {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;
@@ -26,11 +28,11 @@ class b {
             `;
   }
 }
-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: ");
+function f(o) {
+  const { logits: t, labels: e, dy: a } = o.inputs, c = o.backend, r = e.shape[0], i = t.shape[1];
+  s(a.shape, [r], "Error in EfficientScatterSub dy: "), s(t.shape, [r, i], "Error in EfficientScatterSub logits: "), s(e.shape, [r], "Error in EfficientScatterSub labels: ");
   const n = new b(r, i);
-  return c.runWebGPUProgram(n, [e, t, s], "float32");
+  return c.runWebGPUProgram(n, [e, t, a], "float32");
 }
 const l = {
   kernelName: "EfficientScatterSub",

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

@@ -0,0 +1,7 @@
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
+import { SliceAttrs, SliceInputs, TensorInfo } from '@tensorflow/tfjs-core';
+export declare function slice(args: {
+    inputs: SliceInputs;
+    backend: WebGPUBackend;
+    attrs: SliceAttrs;
+}): TensorInfo;