@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/ops/webgpu/normRMSGrad.js

@@ -0,0 +1,225 @@
+import { Fs as e, Ii as t, ii as n } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as r } from "../../utilities/packed.js";
+import { reshape16 as i } from "../reshape16.js";
+import { sum16 as a } from "../sum16.js";
+import { r as o, t as s } from "../../pack16-Ck-spx_F.js";
+import { s as c } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as l } from "../../webgpu_util-B_F3SShA.js";
+import { slice16 as u } from "../slice16.js";
+import { createReduceInfo as d } from "./utils/reductions.js";
+//#region lib/ops/webgpu/normRMSGrad.ts
+var f = class {
+	outputShape;
+	shaderKey = "RMSNormGrad";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = [
+		"x",
+		"gamma",
+		"dy"
+	];
+	uniforms = "reduceSize : i32, batchSize: i32";
+	inputShape;
+	size = !1;
+	rowsPerWorkgroup;
+	packed = !1;
+	outputComponent;
+	constructor(e, t = 4, n = !1) {
+		if (this.packed = n, this.shaderKey = `RMSNormGrad_${t}`, this.rowsPerWorkgroup = t, this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize + e.batchSize / this.rowsPerWorkgroup, e.inSize], this.dispatchLayout = l(this.outputShape), this.dispatch = [
+			e.batchSize / this.rowsPerWorkgroup,
+			1,
+			1
+		], e.batchSize % this.rowsPerWorkgroup !== 0) throw Error(`RMSNormGradProgram: batch size ${e.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`);
+		if (e.inSize > 1024) throw Error(`RMSNormGradProgram: inSize ${e.inSize} exceeds max of 1024`);
+	}
+	getUserCode() {
+		let e = this.workgroupSize[0], t = this.rowsPerWorkgroup, n = `
+          var<workgroup> partials : array<vec2<f32>, ${e}>;
+          var<workgroup> accumulation: array<${this.packed ? "vec2<f32>" : "f32"}, 1024>;
+        `, r = this.packed ? "\n                let X = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                let G  = unpack2x16float(u32(gamma[k]));\n                sum_x2 = fma(X.x, X.x, sum_x2);\n                sum_x2 = fma(X.y, X.y, sum_x2);\n                sum_dygx = fma(DY.x * G.x, X.x, sum_dygx);\n                sum_dygx = fma(DY.y * G.y, X.y, sum_dygx);\n        " : "\n                let X = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                let G  = f32(gamma[k]);\n                sum_x2 = fma(X, X, sum_x2);\n                sum_dygx = fma(DY * G, X, sum_dygx);\n            ", i = this.packed ? "\n                let X  = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                let G  = unpack2x16float(u32(gamma[k]));\n\n                let dyGamma = DY * G;\n                let dx = vec2<f32>(\n                    fma(dyGamma.x, invRMS, -X.x * scale),\n                    fma(dyGamma.y, invRMS, -X.y * scale)\n                );\n\n                result[offset + k] = i32(pack2x16float(dx));\n\n                // dGamma\n                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);\n        " : "\n                let X  = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                let G  = f32(gamma[k]);\n\n                let dyGamma = DY * G;\n                let dx = fma(dyGamma, invRMS, -X * scale);\n\n                result[offset + k] = dx;\n\n                // dGamma\n                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);\n        ", a = this.packed ? "\n                result[outDgBase + k] = i32(pack2x16float(accumulation[k]));\n        " : "\n                result[outDgBase + k] = accumulation[k];\n        ";
+		return `
+          fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+          }
+
+          ${n}
+
+          ${c("index")} {
+            // One workgroup per row (batch).
+            let Length = uniforms.reduceSize;
+            let BatchSize = uniforms.batchSize;
+            for (var k = i32(localId.x); k < Length; k = k + ${e}) {
+                accumulation[k] = ${this.packed ? "vec2<f32>(0.0f)" : "0.0f"};
+            }
+
+            for (var rowOff = 0; rowOff < ${t}; rowOff = rowOff + 1) {
+                let row = i32(workgroupId.x) * ${t} + rowOff;
+                let offset = row * Length;
+                
+                var sum_x2 = 0.0f;
+                var sum_dygx = 0.0f;
+
+                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
+                    ${r}
+                }
+
+                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
+                workgroupBarrier();
+
+                var reduceSize = min(u32(Length), ${e}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.0f / f32(${this.packed ? "Length * 2" : "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 + ${e}) {
+                    ${i}
+                }
+
+                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 + ${e}) {
+                ${a}
+            }
+          }
+        `;
+	}
+}, p = class {
+	outputShape;
+	shaderKey = "RMSNormGrad";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = ["x", "dy"];
+	uniforms = "reduceSize : i32, batchSize: i32";
+	inputShape;
+	size = !1;
+	packed = !1;
+	outputComponent;
+	constructor(e, t = !1) {
+		this.packed = t, this.shaderKey = "RMSNormGrad_NoGamma", this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize, e.inSize], this.dispatchLayout = l(this.outputShape), this.dispatch = [
+			e.batchSize,
+			1,
+			1
+		];
+	}
+	getUserCode() {
+		let e = this.workgroupSize[0], t = `
+          var<workgroup> partials : array<vec2<f32>, ${e}>;
+        `, n = this.packed ? "\n                let X = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n                sum_x2 = fma(X.x, X.x, sum_x2);\n                sum_x2 = fma(X.y, X.y, sum_x2);\n                sum_dygx = fma(DY.x, X.x, sum_dygx);\n                sum_dygx = fma(DY.y, X.y, sum_dygx);\n        " : "\n                let X = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n                sum_x2 = fma(X, X, sum_x2);\n                sum_dygx = fma(DY, X, sum_dygx);\n            ", r = this.packed ? "\n                let X  = unpack2x16float(u32(x[offset + k]));\n                let DY = unpack2x16float(u32(dy[offset + k]));\n\n                let dx = vec2<f32>(\n                    fma(DY.x, invRMS, -X.x * scale),\n                    fma(DY.y, invRMS, -X.y * scale)\n                );\n\n                result[offset + k] = i32(pack2x16float(dx));\n        " : "\n                let X  = f32(x[offset + k]);\n                let DY = f32(dy[offset + k]);\n\n                let dx = fma(DY, invRMS, -X * scale);\n\n                result[offset + k] = dx;\n        ";
+		return `
+          fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+          }
+
+          ${t}
+
+          ${c("index")} {
+            // One workgroup per row (batch).
+            let Length = uniforms.reduceSize;
+            let BatchSize = uniforms.batchSize;
+
+                let row = i32(workgroupId.x);
+                let offset = row * Length;
+                
+                var sum_x2 = 0.0f;
+                var sum_dygx = 0.0f;
+
+                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
+                    ${n}
+                }
+
+                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
+                workgroupBarrier();
+
+                var reduceSize = min(u32(Length), ${e}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.0f / f32(${this.packed ? "Length * 2" : "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 + ${e}) {
+                    ${r}
+                }
+          }
+        `;
+	}
+};
+function m(t) {
+	let { dy: c, x: l, gamma: m } = t.inputs;
+	e(l.shape, c.shape, "Error in RMSNormGrad dy: ");
+	let h = r(l), g = m ? r(m) : !1, _ = r(c), v = h || g || _, y = !v || h ? l : s(l), b = !v || g || !m ? m : s(m), x = !v || _ ? c : s(c);
+	b && e(b.shape, [y.shape[y.shape.length - 1]], "Error in RMSNormGrad gamma: ");
+	let S = t.backend, C = d(b ? [
+		y,
+		b,
+		x
+	] : [y, x], -1), w = m ? new f(C, 4, v) : new p(C, v), T = [{
+		type: "int32",
+		data: [w.inputShape[1]]
+	}, {
+		type: "int32",
+		data: [w.inputShape[0]]
+	}];
+	if (C.inSize > 1024) throw Error(`rmsNormGradGPU: inSize ${C.inSize} exceeds max of 1024`);
+	let E = S.runWebGPUProgram(w, b ? [
+		y,
+		b,
+		x
+	] : [y, x], v ? "packedF16" : "float32", T);
+	v && !h && y.dispose(), v && !g && b && b.dispose(), v && !_ && x.dispose();
+	let D = n().makeTensorFromTensorInfo(E);
+	if (m) {
+		let e = u(D, [0, 0], [C.batchSize, C.inSize]), t = u(D, [C.batchSize, 0], [C.batchSize / 4, C.inSize]);
+		D.dispose();
+		let n = i(e, l.shape);
+		e.dispose();
+		let r = a(t, [0]);
+		return t.dispose(), [n, !v || g ? r : o(r)];
+	} else {
+		let e = i(D, l.shape);
+		return D.dispose(), [e];
+	}
+}
+t({
+	kernelName: "RMSNormGrad",
+	backendName: "webgpu",
+	kernelFunc: m
+});
+//#endregion

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

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

package/dist/ops/webgpu/pack16.js

@@ -0,0 +1,21 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+import t from "./pack16_program.js";
+//#region lib/ops/webgpu/pack16.ts
+function n(e) {
+	let { x: n } = e.inputs, { scaling: r, padding: i } = e.attrs, a = e.backend;
+	if (n.shape[n.shape.length - 1] % 2 != 0) throw Error("Last dimension of input tensor must be even to use Pack16.");
+	e.attrs && (e.attrs.originalShape = n.shape);
+	let o = new t(n.shape, i), s = r !== 1;
+	s && o.useScaling();
+	let c = [{
+		type: "float32",
+		data: [r]
+	}];
+	return a.runWebGPUProgram(o, [n], "packedF16", s ? c : void 0);
+}
+e({
+	kernelName: "Pack16",
+	backendName: "webgpu",
+	kernelFunc: n
+});
+//#endregion

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,93 @@
+import { s as e } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as t, i as n } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/pack16_program.ts
+var r = class {
+	outputShape;
+	shaderKey = "Pack16";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = ["x"];
+	uniforms;
+	size = !0;
+	outputComponent = 4;
+	scaling = !1;
+	padding = 0;
+	constructor(e, r = 0) {
+		if (e[e.length - 1] % 2 != 0 && r === 0) throw Error("Last dimension of output shape must be even to use Pack16.");
+		if (r % 4 != 0) throw Error("Padding must be a multiple of 4 to use Pack16.");
+		if (this.outputShape = [...e.slice(0, -1), e[e.length - 1]], r > 0) {
+			this.shaderKey += `_Padded${r}`, this.padding = r;
+			for (let e = this.outputShape.length - 2; e < this.outputShape.length; e++) this.outputShape[e] % this.padding !== 0 && (this.outputShape[e] += this.padding - this.outputShape[e] % 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 = t(this.outputShape), this.dispatch = n(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) {
+			let t = this.outputShape.length;
+			return `
+                ${e("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 ? `
+                ${e("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;
+                    }
+                }` : `
+                ${e("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;
+                    }
+                }`;
+	}
+};
+//#endregion
+export { r as default };

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

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

package/dist/ops/webgpu/qkv.js

@@ -0,0 +1,64 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import { t as r } from "../../matMul16-BNfZSnNM.js";
+import { slice16 as i } from "../slice16.js";
+//#region lib/ops/webgpu/qkv.ts
+function a(t) {
+	let { x: a, kernel: o } = t.inputs, { heads: s } = t.attrs, c = a.shape[0], l = a.shape[1], u = a.shape[2], d = n(a);
+	if (e(o.shape, [d ? u * 2 : u, 3 * u], "Error in QKV: "), u % s !== 0) throw Error(`Channel dimension ${u} must be divisible by number of heads ${s} in QKV.`);
+	let f = r(a, o, !1, !1, {
+		forceOutputShape: [
+			c,
+			l,
+			3 * s,
+			u / s
+		],
+		perm: [
+			0,
+			2,
+			1,
+			3
+		]
+	}), p = [
+		i(f, [
+			0,
+			0,
+			0,
+			0
+		], [
+			c,
+			s,
+			l,
+			u / s
+		]),
+		i(f, [
+			0,
+			s,
+			0,
+			0
+		], [
+			c,
+			s,
+			l,
+			u / s
+		]),
+		i(f, [
+			0,
+			2 * s,
+			0,
+			0
+		], [
+			c,
+			s,
+			l,
+			u / s
+		])
+	];
+	return f.dispose(), p;
+}
+t({
+	kernelName: "QKV",
+	backendName: "webgpu",
+	kernelFunc: a
+});
+//#endregion

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

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

package/dist/ops/webgpu/rope.js

@@ -0,0 +1,163 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import { s as r } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as i, i as a } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/rope.ts
+var o = class {
+	variableNames = [
+		"x",
+		"sin",
+		"cos"
+	];
+	outputShape;
+	shaderKey = "Rope";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "pastLen: i32";
+	constructor(e, t, n, r) {
+		this.shaderKey = `Rope_${r}`, this.outputShape = [
+			e,
+			t,
+			n,
+			r
+		], this.dispatchLayout = i(this.outputShape), this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		let e = this.outputShape[3];
+		return `
+        ${r("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];
+
+                let rotaryDim = ${e};
+
+                var outVal = 0.0;
+
+                let xIdx = b * uniforms.outShapeStrides[0] +
+                    h * uniforms.outShapeStrides[1] +
+                    t * uniforms.outShapeStrides[2] +
+                    d;
+
+                if (d < rotaryDim) {
+                    let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d / 2;
+                    let cos = cos[idx];
+                    let sin = sin[idx];
+
+                    let ownX = x[xIdx] * cos;
+                    var evenOdd = 0.0;
+
+                    if (d % 2 == 0) {
+                        // even index
+                        evenOdd = -x[xIdx + 1];
+                    } else {
+                        // odd index
+                        evenOdd = x[xIdx - 1];
+                    }
+
+                    outVal = fma(evenOdd, sin, ownX);
+                } else {
+                    // pass through for non-rotary dims
+                    outVal = x[xIdx];
+                }
+
+                setOutputAtIndex(index, outVal);
+            }
+        }
+        `;
+	}
+}, s = class {
+	variableNames = [
+		"x",
+		"sin",
+		"cos"
+	];
+	outputShape;
+	shaderKey = "Rope";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "pastLen: i32";
+	constructor(e, t, n, r) {
+		this.shaderKey = `Rope_${r}`, this.outputShape = [
+			e,
+			t,
+			n,
+			r / 2
+		], this.dispatchLayout = i(this.outputShape), this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		return `
+        ${r("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 c(t) {
+	let { x: r } = t.inputs, { pastLen: i, negSin: a, ropeCache: c } = t.attrs, l = t.backend, u = n(r), d = r.shape[0], f = r.shape[1], p = r.shape[2], m = u ? r.shape[3] * 2 : r.shape[3], h = a ? c.getNegSin() : c.getSin(), g = c.getCos();
+	if (e(h.shape, g.shape, "Error in Rope: "), h.shape[0] < p + i) throw Error(`Sin tensor shape ${h.shape} is not compatible with seqLength ${p} and pastLen ${i}.`);
+	if (h.shape[1] * 2 < m) throw Error(`Sin tensor shape ${h.shape} is not compatible with feature dimension ${m}.`);
+	if (h.shape.length !== 3) throw Error(`Sin tensor must be 3-dimensional, but got shape ${h.shape}.`);
+	let _ = u ? new s(d, f, p, m) : new o(d, f, p, m), v = [{
+		type: "int32",
+		data: [i]
+	}], y = u ? "packedF16" : r.dtype;
+	return l.runWebGPUProgram(_, [
+		r,
+		h,
+		g
+	], y, v);
+}
+t({
+	kernelName: "Rope",
+	backendName: "webgpu",
+	kernelFunc: c
+});
+//#endregion

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

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

package/dist/ops/webgpu/scatterSub.js

@@ -0,0 +1,53 @@
+import { Fs as e, Ii as t } from "../../dist-BewPQWjc.js";
+import { s as n } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as r, i } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/scatterSub.ts
+var a = class {
+	variableNames = [
+		"labels",
+		"softmaxProbs",
+		"dy"
+	];
+	outputShape;
+	shaderKey = "ScatterSub";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	constructor(e, t) {
+		this.outputShape = [e, t], this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		return `
+            ${n("index")} {
+                if (index < uniforms.size) {
+                    let coords = getCoordsFromIndex(index); // [batch, depth]
+                    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 o(t) {
+	let { logits: n, labels: r, dy: i } = t.inputs, o = t.backend, s = r.shape[0], c = n.shape[1];
+	e(i.shape, [s], "Error in EfficientScatterSub dy: "), e(n.shape, [s, c], "Error in EfficientScatterSub logits: "), e(r.shape, [s], "Error in EfficientScatterSub labels: ");
+	let l = new a(s, c);
+	return o.runWebGPUProgram(l, [
+		r,
+		n,
+		i
+	], "float32");
+}
+t({
+	kernelName: "EfficientScatterSub",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+//#endregion

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;