@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/ops/webgpu/slice16.js

@@ -0,0 +1,74 @@
+import { Ii as e, oc as t } from "../../dist-BewPQWjc.js";
+import { a as n, t as r } from "../../slice_util-Dgb3ANWI.js";
+import { i, o as a, s as o } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as s, i as c } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/slice16.ts
+function l(e) {
+	switch (e) {
+		case 1: return "1D";
+		case 2: return "2D";
+		case 3: return "3D";
+		case 4: return "4D";
+	}
+	return "";
+}
+var u = class {
+	variableNames = ["source"];
+	uniforms;
+	outputShape;
+	shaderKey;
+	rank;
+	dispatchLayout;
+	dispatch;
+	workPerThread = 1;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	start;
+	size = !0;
+	constructor(e, t) {
+		this.outputShape = t, this.rank = t.length, this.dispatchLayout = s(this.outputShape), this.dispatch = c(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			this.workPerThread,
+			1,
+			1
+		]), this.start = e, this.uniforms = `start : ${i(e.length)}, `, this.shaderKey = "slice";
+	}
+	getUserCode() {
+		let e = i(this.rank), t;
+		return t = this.start.length === 1 ? this.outputShape.map(() => "sourceLoc = uniforms.start + coords;") : this.outputShape.map((e, t) => `sourceLoc.${d[t]} = uniforms.start.${a(t)} + coords.${d[t]};`), `
+      ${o("index")} {
+        if (index < uniforms.size) {
+          var sourceLoc : ${e};
+          let coords = getCoordsFromIndex(index);
+          ${t.join("\n")}
+          result[index] = source[getIndexFromCoords${l(this.rank)}(sourceLoc, uniforms.sourceShape)];
+        }
+      }
+    `;
+	}
+}, d = [
+	"x",
+	"y",
+	"z",
+	"w",
+	"u",
+	"v"
+];
+function f(e) {
+	let { inputs: i, backend: a, attrs: o } = e, { x: s } = i, { begin: c, size: l } = o, [d, f] = n(s, c, l);
+	if (r(s, d, f), t(f) === 0) return a.makeTensorInfo(f, s.dtype, []);
+	let p = new u(d, f), m = [{
+		type: "int32",
+		data: d
+	}];
+	return a.runWebGPUProgram(p, [s], s.dtype, m);
+}
+e({
+	kernelName: "Slice16",
+	backendName: "webgpu",
+	kernelFunc: f
+});
+//#endregion
+export { f as slice };

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

@@ -0,0 +1,17 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+export {};

package/dist/ops/webgpu/softmax16.js

@@ -0,0 +1,18 @@
+import { Ii as e, br as t, ii as n, oc as r } from "../../dist-BewPQWjc.js";
+import { reshape16 as i } from "../reshape16.js";
+import a from "./utils/deviceInfo.js";
+import o from "./softmax16_subgroup_program.js";
+import s from "./softmax16_program.js";
+//#region lib/ops/webgpu/softmax16.ts
+function c(e) {
+	let { inputs: c, backend: l, attrs: u } = e, { logits: d } = c, { dim: f } = u, p = l.subgroupMinSize, m = l.subgroupMaxSize, h = a(l).subgroupsSupported, g = t(d, [r(d.shape) / d.shape[f], d.shape[f]]), _ = h ? new o(g.shape, p, m) : new s(g.shape), v = l.runWebGPUProgram(_, [g], "packedF16");
+	g.dispose();
+	let y = n().makeTensorFromTensorInfo(v), b = i(y, d.shape);
+	return y.dispose(), b;
+}
+e({
+	kernelName: "Softmax16",
+	backendName: "webgpu",
+	kernelFunc: c
+});
+//#endregion

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

@@ -0,0 +1,13 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class SoftmaxProgram implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    constructor(outputShape: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/softmax16_program.js

@@ -0,0 +1,89 @@
+import { s as e } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as t } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/softmax16_program.ts
+var n = class {
+	variableNames = ["logits"];
+	outputShape;
+	shaderKey;
+	dispatchLayout;
+	dispatch;
+	workgroupSize;
+	constructor(e) {
+		this.outputShape = e, this.dispatchLayout = t(this.outputShape), this.dispatch = [
+			this.outputShape[0],
+			1,
+			1
+		], this.outputShape[1] >= 4096 ? this.workgroupSize = [
+			256,
+			1,
+			1
+		] : this.outputShape[1] < 64 ? this.workgroupSize = [
+			32,
+			1,
+			1
+		] : this.workgroupSize = [
+			64,
+			1,
+			1
+		], this.shaderKey = "softmax16";
+	}
+	getUserCode() {
+		return `
+        var<workgroup> buf : array<f32, ${this.workgroupSize[0]}>;
+        const blockSize = ${this.workgroupSize[0]};
+        ${e("index")} {
+            let row = index / blockSize;
+            let tid = i32(localId.x);
+            let cols = uniforms.outShape[1];
+            let rowIdx = row * cols;
+
+            var threadMax = -3.402823e+38f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                threadMax = max(threadMax, max(value.x, value.y));
+            }
+            buf[tid] = threadMax;
+            workgroupBarrier();
+
+            for (var currSize = blockSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+                if (tid < currSize) {
+                    buf[tid] = max(buf[tid], buf[tid + currSize]);
+                }
+                workgroupBarrier();
+            }
+
+            let rowMaxShared: f32 = buf[0];
+            workgroupBarrier();
+
+            var threadSum = 0.0f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let subExp = exp(value.x - rowMaxShared);
+                threadSum += subExp;
+                let subExpY = exp(value.y - rowMaxShared);
+                threadSum += subExpY;
+            }
+            buf[tid] = threadSum;
+            workgroupBarrier();
+
+            for (var currSize = blockSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+                if (tid < currSize) {
+                    buf[tid] = buf[tid] + buf[tid + currSize];
+                }
+                workgroupBarrier();
+            }
+
+            let rowSumShared: f32 = buf[0];
+
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let value1: f32 = exp(value.x - rowMaxShared) / rowSumShared;
+                let value2: f32 = exp(value.y - rowMaxShared) / rowSumShared;
+                result[rowIdx + col] = i32(pack2x16float(vec2<f32>(value1, value2)));
+            }
+        }
+      `;
+	}
+};
+//#endregion
+export { n as default };

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

@@ -0,0 +1,17 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class SoftmaxSubgroupProgram implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    minSubgroupSize: number;
+    maxSubgroupSize: number;
+    subgroups: boolean;
+    subgroupBuiltins: boolean;
+    constructor(outputShape: number[], minSubgroupSize: number, maxSubgroupSize: number);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/softmax16_subgroup_program.js

@@ -0,0 +1,70 @@
+import { s as e } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as t } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/softmax16_subgroup_program.ts
+var n = class {
+	variableNames = ["logits"];
+	outputShape;
+	shaderKey;
+	dispatchLayout;
+	dispatch;
+	workgroupSize;
+	minSubgroupSize;
+	maxSubgroupSize;
+	subgroups = !0;
+	subgroupBuiltins = !1;
+	constructor(e, n, r) {
+		this.minSubgroupSize = n, this.maxSubgroupSize = r, this.outputShape = e, this.dispatchLayout = t(this.outputShape), this.dispatch = [
+			this.outputShape[0],
+			1,
+			1
+		], this.workgroupSize = [
+			Math.min(64, r),
+			1,
+			1
+		], n !== r && (this.subgroupBuiltins = !0, this.workgroupSize = [
+			64,
+			1,
+			1
+		]), this.shaderKey = "softmax16subgroup";
+	}
+	getUserCode() {
+		let t = this.maxSubgroupSize !== this.minSubgroupSize;
+		return `
+        ${t ? `var<workgroup> bestValues : array<f32, ${this.workgroupSize[0]}>;` : ""}
+        const blockSize = ${this.workgroupSize[0]};
+        ${e("index")} {
+            let row = index / blockSize;
+            let tid = i32(localId.x);
+            let cols = uniforms.outShape[1];
+            let rowIdx = row * cols;
+
+            var threadMax = -3.402823e+38f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                threadMax = max(threadMax, max(value.x, value.y));
+            }
+
+            threadMax = subgroupMax(threadMax);
+            ${t ? "\n                let lane = localId.x % subgroupSize;\n                if (lane == 0) {\n                    bestValues[localId.x / subgroupSize] = threadMax;\n                }\n                workgroupBarrier();\n                let numSubgroups = blockSize / subgroupSize;\n                threadMax = select(-3.402823e+38f, bestValues[lane], lane < numSubgroups);\n                threadMax = subgroupMax(threadMax);\n                workgroupBarrier();    \n            " : ""}
+
+            var threadSum = 0.0f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let subExp = exp(value - threadMax);
+                threadSum += subExp.x + subExp.y;
+            }
+
+            threadSum = subgroupAdd(threadSum);
+            ${t ? "\n                if (lane == 0) {\n                    bestValues[localId.x / subgroupSize] = threadSum;\n                }\n                workgroupBarrier();\n                threadSum = select(0.0f, bestValues[lane], lane < numSubgroups);\n                threadSum = subgroupAdd(threadSum);    \n            " : ""}
+
+            for (var col = tid; col < cols; col += ${t ? "i32(subgroupSize)" : "blockSize"}) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let valuePair: vec2<f32> = exp(value - threadMax) / threadSum;
+                result[rowIdx + col] = i32(pack2x16float(valuePair));
+            }
+        }
+      `;
+	}
+};
+//#endregion
+export { n as default };

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

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

package/dist/ops/webgpu/softmax16grad.js

@@ -0,0 +1,31 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as t } from "../../utilities/packed.js";
+import { ReduceProgram as n, createReduceInfo as r, reduce as i } from "./utils/reductions.js";
+import a from "./utils/deviceInfo.js";
+//#region lib/ops/webgpu/softmax16grad.ts
+var o = class extends n {
+	constructor(e, t) {
+		super(e, t, {
+			reductionOp: "sum",
+			elementwise: !0
+		}, !0), this.shaderKey = "SoftmaxGrad16", this.variableNames = ["dy", "softmaxOutput"], this.variableComponents = [1, 1];
+	}
+	getReadSnippet() {
+		return "\n            let d: vec2<f32> = unpack2x16float(u32(dy[index]));\n            let l: vec2<f32> = unpack2x16float(u32(softmaxOutput[index]));\n            return d * l;\n        ";
+	}
+	getWriteSnippet() {
+		return "\n            let d: vec2<f32> = unpack2x16float(u32(dy[offset + k]));\n            let l: vec2<f32> = unpack2x16float(u32(softmaxOutput[offset + k]));\n            let outVal = l * (d - bestValue);\n            result[offset + k] = i32(pack2x16float(outVal));\n        ";
+	}
+};
+function s(e) {
+	let { dy: n, softmaxOutput: s } = e.inputs, c = e.backend, l = a(c), u = t(n), d = t(s);
+	if (!(u && d)) throw Error("softmaxGradGPU: dy and softmaxOutput must be packed tensors");
+	let f = [n, s];
+	return i(new o(l, r(f, -1)), f, c);
+}
+e({
+	kernelName: "Softmax16Grad",
+	backendName: "webgpu",
+	kernelFunc: s
+});
+//#endregion

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

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

package/dist/ops/webgpu/sub16.js

@@ -0,0 +1,14 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+import { t } from "../../binary_op_util-CrYk9LXL.js";
+import { BinaryOpProgram as n } from "./utils/binary_op.js";
+//#region lib/ops/webgpu/sub16.ts
+function r(e) {
+	let { a: r, b: i } = e.inputs, a = e.backend, o = new n(t.SUB, r.shape, i.shape);
+	return a.runWebGPUProgram(o, [r, i], "packedF16");
+}
+e({
+	kernelName: "Sub16",
+	backendName: "webgpu",
+	kernelFunc: r
+});
+//#endregion

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

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

package/dist/ops/webgpu/sum16.js

@@ -0,0 +1,29 @@
+import { Dn as e, En as t, Ii as n, mn as r, nc as i } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as a } from "../../utilities/packed.js";
+import { transpose16 as o } from "../transpose16.js";
+import { ReduceProgram as s, createReduceInfo as c, reduce as l } from "./utils/reductions.js";
+import u from "./utils/deviceInfo.js";
+//#region lib/ops/webgpu/sum16.ts
+var d = class extends s {
+	shaderKey = "sum16";
+	constructor(e, t, n) {
+		super(e, t, {
+			reductionOp: "sum",
+			elementwise: !1
+		}, n), n && (this.shaderKey += "_packed");
+	}
+};
+function f(n) {
+	let { x: s } = n.inputs, { axis: f, keepDims: p } = n.attrs, m = n.backend, h = [], g = u(m), _ = a(s);
+	if (!_) return r(s, f, p);
+	let v = i(f ?? -1, s.shape), y = t(v, s.shape.length), b = s;
+	y != null && (b = o(s, y), v = e(v.length, b.shape.length), h.push(b));
+	let x = l(new d(g, c([b], -1), _), [b], m);
+	return h.forEach((e) => e.dispose()), x;
+}
+n({
+	kernelName: "Sum16",
+	backendName: "webgpu",
+	kernelFunc: f
+});
+//#endregion

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

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

package/dist/ops/webgpu/transpose16.js

@@ -0,0 +1,37 @@
+import { Ii as e, ii as t, w as n } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as r } from "../../utilities/packed.js";
+import { reshape16 as i } from "../reshape16.js";
+import a from "./transpose16_shared_program.js";
+import o from "./transpose16_program.js";
+//#region lib/ops/webgpu/transpose16.ts
+function s(e) {
+	let { inputs: s, attrs: c } = e, { x: l } = s, { perm: u } = c, d = e.backend, f = r(l);
+	if (f && u[u.length - 1] !== l.shape.length - 1) {
+		let e = l.shape.length, n = e === 4 ? u.map((e) => e - 1).filter((e) => e >= 0) : u, r = e === 4 ? i(l, [
+			l.shape[0] * l.shape[1],
+			l.shape[2],
+			l.shape[3]
+		]) : l, o = new a(r.shape, n), s = d.runWebGPUProgram(o, [r], "packedF16");
+		if (e === 4) {
+			r.dispose();
+			let e = t().makeTensorFromTensorInfo(s), n = i(e, [
+				l.shape[0],
+				l.shape[1],
+				s.shape[1],
+				s.shape[2]
+			]);
+			return e.dispose(), n;
+		}
+		return s;
+	}
+	if (f) {
+		let e = new o(l.shape, u);
+		return d.runWebGPUProgram(e, [l], "packedF16");
+	} else return n(l, u);
+}
+e({
+	kernelName: "Transpose16",
+	backendName: "webgpu",
+	kernelFunc: s
+});
+//#endregion

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

@@ -0,0 +1,16 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class TransposeProgram16 implements WebGPUProgram {
+    variableNames: string[];
+    shaderKey: string;
+    outputShape: number[];
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workPerThread: number;
+    workgroupSize: [number, number, number];
+    newDim: number[];
+    size: boolean;
+    constructor(aShape: number[], newDim: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/transpose16_program.js

@@ -0,0 +1,51 @@
+import { i as e, o as t, s as n } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as r, i } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/transpose16_program.ts
+function a(e) {
+	let n = e.length;
+	if (n > 6) throw Error(`Transpose for rank ${n} is not yet supported`);
+	let r = Array(n);
+	for (let n = 0; n < e.length; n++) r[e[n]] = `coords.${t(n)}`;
+	return r.join();
+}
+var o = class {
+	variableNames = ["A"];
+	shaderKey;
+	outputShape;
+	dispatchLayout;
+	dispatch;
+	workPerThread = 1;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	newDim;
+	size = !0;
+	constructor(e, t) {
+		let n = Array(e.length);
+		for (let r = 0; r < n.length; r++) n[r] = e[t[r]];
+		this.outputShape = n, this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			this.workPerThread,
+			1,
+			1
+		]), this.newDim = t, this.shaderKey = `transpose16_${t}`;
+	}
+	getUserCode() {
+		let t = e(this.outputShape.length), r = a(this.newDim);
+		return `
+      ${n("index")} {
+        for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
+          let flatIndex = index * ${this.workPerThread} + i;
+          if(flatIndex < uniforms.size) {
+            let coords = getCoordsFromIndex(flatIndex);
+            result[flatIndex] = A[getIndexFromCoords${this.outputShape.length}D(
+              ${t}(${r}), uniforms.aShape)];
+          }
+        }
+      }
+    `;
+	}
+};
+//#endregion
+export { o as default };

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

@@ -0,0 +1,15 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
+export default class TransposeSharedProgram16 implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+        y: number[];
+        z?: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    constructor(aShape: number[], newDim: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/transpose16_shared_program.js

@@ -0,0 +1,79 @@
+import { Ps as e } from "../../dist-BewPQWjc.js";
+import { s as t } from "../../webgpu_program-WOyIVMlZ.js";
+import { i as n } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/transpose16_shared_program.ts
+var r = class {
+	variableNames = ["A"];
+	outputShape;
+	shaderKey;
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		8,
+		8,
+		1
+	];
+	constructor(e, t) {
+		let r = e.length, i = Array(r), a = e.slice();
+		a[a.length - 1] *= 2;
+		for (let e = 0; e < i.length; e++) i[e] = a[t[e]];
+		i[i.length - 1] /= 2, this.outputShape = i, this.dispatchLayout = r === 2 ? {
+			x: [0],
+			y: [1]
+		} : {
+			x: [1],
+			y: [2],
+			z: [0]
+		}, this.dispatch = n(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			2,
+			1,
+			1
+		]), this.shaderKey = `transposeShared16_${r}`;
+	}
+	getUserCode() {
+		let n = this.outputShape.length;
+		e(this.workgroupSize[0] === this.workgroupSize[1], () => `Must be a square tile, current tile shape is ${this.workgroupSize[0]} x ${this.workgroupSize[1]}`);
+		let r = this.workgroupSize[0] * 2;
+		return `
+      var<workgroup> tile : array<array<f32, ${r + 1}>, ${r}>;
+      ${t()} {
+        var x = i32(workgroupId.x) * ${r / 2} + i32(localId.x);
+        var y = i32(workgroupId.y) * ${r} + i32(localId.y);
+        let batch = ${n === 3 ? "i32(workgroupId.z)" : "0"};
+        let batchOffsetA = ${n === 3 ? "batch * uniforms.aShapeStrides[0]" : "0"};
+        let batchOffsetOut = ${n === 3 ? "batch * uniforms.outShapeStrides[0]" : "0"};
+
+        let inputWidth = uniforms.outShape[${n === 3 ? "1" : "0"}] / 2; // Output height
+        let inputHeight = uniforms.outShape[${n === 3 ? "2" : "1"}] * 2; // Output width
+        if (x < inputWidth && y < inputHeight) {
+            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
+            tile[localId.y][localId.x * 2] = unpackedA.x;
+            tile[localId.y][localId.x * 2 + 1] = unpackedA.y;
+        }
+        // Second load to cover the tile
+        y = y + ${this.workgroupSize[0]};
+        if (x < inputWidth && y < inputHeight) {
+            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
+            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2] = unpackedA.x;
+            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2 + 1] = unpackedA.y;
+        }
+        workgroupBarrier();
+
+        let outputWidth = uniforms.outShape[${n === 3 ? "2" : "1"}]; // Output width
+        let outputHeight = uniforms.outShape[${n === 3 ? "1" : "0"}] * 2; // Output height
+        x = i32(workgroupId.y) * ${r / 2} + i32(localId.x);
+        y = i32(workgroupId.x) * ${r} + i32(localId.y);
+        if (x < outputWidth && y < outputHeight) {
+          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y], tile[localId.x * 2 + 1][localId.y])));
+        }
+        // Second store to cover the tile
+        y = y + ${this.workgroupSize[0]};
+        if (x < outputWidth && y < outputHeight) {
+          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y + ${this.workgroupSize[0]}], tile[localId.x * 2 + 1][localId.y + ${this.workgroupSize[0]}])));
+        }
+      }
+    `;
+	}
+};
+//#endregion
+export { r as default };

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

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

package/dist/ops/webgpu/unpack16.js

@@ -0,0 +1,60 @@
+import { Ii as e } from "../../dist-BewPQWjc.js";
+import { s as t } from "../../webgpu_program-WOyIVMlZ.js";
+import { c as n, i as r } from "../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/unpack16.ts
+var i = class {
+	outputShape;
+	shaderKey = "Unpack16";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = ["x"];
+	size = !0;
+	uniforms;
+	outputComponent = 4;
+	variableComponents = [2];
+	scaling = !1;
+	constructor(e) {
+		this.outputShape = [...e.slice(0, -1), e[e.length - 1] * 2], this.dispatchLayout = n(this.outputShape), this.dispatch = r(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			4,
+			1,
+			1
+		]);
+	}
+	useScaling() {
+		this.shaderKey += "_Scaled", this.uniforms = "scaling : f32,", this.scaling = !0;
+	}
+	getUserCode() {
+		return `
+                ${t("index")} {
+                    let outIndex = index;
+                    if (outIndex < uniforms.size) {
+                        let xvec2 = x[index];
+                        let v1 = vec4<f32>(
+                            unpack2x16float(u32(xvec2.x)),
+                            unpack2x16float(u32(xvec2.y))
+                        ) ${this.scaling ? "* uniforms.scaling" : ""};
+                        result[outIndex] = v1;
+                    }
+                }`;
+	}
+};
+function a(e) {
+	let { x: t } = e.inputs, { scaling: n } = e.attrs, r = e.backend, a = new i(t.shape), o = n !== 1;
+	o && a.useScaling();
+	let s = [{
+		type: "float32",
+		data: [1 / n]
+	}];
+	return r.runWebGPUProgram(a, [t], "float32", o ? s : void 0);
+}
+e({
+	kernelName: "Unpack16",
+	backendName: "webgpu",
+	kernelFunc: a
+});
+//#endregion