@genai-fi/nanogpt 0.20.0 → 0.20.1

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

@@ -0,0 +1,35 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+import { BinaryOpType } from '@tensorflow/tfjs-backend-webgpu/dist/binary_op_util';
+export { BinaryOpType };
+export declare class BinaryOpProgram implements WebGPUProgram {
+    dispatch: [number, number, number];
+    dispatchLayout: {
+        x: number[];
+    };
+    outputComponent: number;
+    op: BinaryOpType;
+    outputShape: number[];
+    shaderKey: string;
+    size: boolean;
+    variableNames: string[];
+    workgroupSize: [number, number, number];
+    variableComponents: number[];
+    constructor(op: BinaryOpType, aShape: number[], bShape: number[]);
+    getUserCode(): string;
+}
+export declare class BinaryOpScalarProgram implements WebGPUProgram {
+    dispatch: [number, number, number];
+    dispatchLayout: {
+        x: number[];
+    };
+    outputComponent: number;
+    op: BinaryOpType;
+    outputShape: number[];
+    shaderKey: string;
+    size: boolean;
+    variableNames: string[];
+    workgroupSize: [number, number, number];
+    variableComponents: number[];
+    constructor(op: BinaryOpType, aShape: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/utils/binary_op.js

@@ -0,0 +1,141 @@
+import { In as e, Js 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";
+import { n as a, t as o } from "../../../binary_op_util-CrYk9LXL.js";
+//#region lib/ops/webgpu/utils/binary_op.ts
+var s = class {
+	dispatch;
+	dispatchLayout;
+	outputComponent;
+	op;
+	outputShape;
+	shaderKey;
+	size = !0;
+	variableNames = ["A", "B"];
+	workgroupSize;
+	variableComponents;
+	constructor(n, a, o) {
+		this.outputShape = e(a, o), this.dispatchLayout = r(this.outputShape), this.op = n;
+		let s = a.length > 0 && a[a.length - 1] % 4 == 0, c = o.length > 0 && o[o.length - 1] % 4 == 0;
+		if (s && c) this.outputComponent = 4, this.variableComponents = [4, 4];
+		else if (s && (t(o) || o[o.length - 1] === 1) || c && (t(a) || a[a.length - 1] === 1)) throw Error("Cannot broadcast 16-bit float binary ops with mixed vector sizes");
+		else throw Error("16-bit float binary ops require inner dimension to be multiple of 4");
+		this.shaderKey = `binary_${n}_${this.variableComponents}`, this.workgroupSize = [
+			128,
+			1,
+			1
+		], this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			this.outputComponent,
+			1,
+			1
+		]);
+	}
+	getUserCode() {
+		let e = this.outputComponent === 4 ? "vec4<f32>" : "f32";
+		return `
+       ${`
+    fn binaryOperation(a : ${e}, b : ${e}) -> ${e} {
+      ${a(this.op, this.outputComponent === 4)}
+    };
+    `}
+       ${n("index")} {
+         if (index < uniforms.size) {
+            let a = A[index];
+            let b = B[index];
+
+            let v4a1 = vec4<f32>(
+                unpack2x16float(u32(a.x)),
+                unpack2x16float(u32(a.y))
+            );
+            let v4a2 = vec4<f32>(
+                unpack2x16float(u32(a.z)),
+                unpack2x16float(u32(a.w))
+            );
+            let v4b1 = vec4<f32>(
+                unpack2x16float(u32(b.x)),
+                unpack2x16float(u32(b.y))
+            );
+            let v4b2 = vec4<f32>(
+                unpack2x16float(u32(b.z)),
+                unpack2x16float(u32(b.w))
+            );
+
+            let v4res1 = binaryOperation(v4a1, v4b1);
+            let v4res2 = binaryOperation(v4a2, v4b2);
+
+            let res = vec4<i32>(
+                i32(pack2x16float(v4res1.xy)),
+                i32(pack2x16float(v4res1.zw)),
+                i32(pack2x16float(v4res2.xy)),
+                i32(pack2x16float(v4res2.zw))
+            );
+
+            result[index] = res;
+         }
+       }
+       `;
+	}
+}, c = class {
+	dispatch;
+	dispatchLayout;
+	outputComponent;
+	op;
+	outputShape;
+	shaderKey;
+	size = !0;
+	variableNames = ["A", "B"];
+	workgroupSize;
+	variableComponents;
+	constructor(e, t) {
+		if (this.outputShape = t, this.dispatchLayout = r(this.outputShape), this.op = e, t.length > 0 && t[t.length - 1] % 4 == 0) this.outputComponent = 4, this.variableComponents = [4, 1];
+		else throw Error("16-bit float binary ops require inner dimension to be multiple of 4");
+		this.shaderKey = `binary_scal_${e}_${this.variableComponents}`, this.workgroupSize = [
+			128,
+			1,
+			1
+		], this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			this.outputComponent,
+			1,
+			1
+		]);
+	}
+	getUserCode() {
+		let e = this.outputComponent === 4 ? "vec4<f32>" : "f32";
+		return `
+       ${`
+    fn binaryOperation(a : ${e}, b : ${e}) -> ${e} {
+      ${a(this.op, this.outputComponent === 4)}
+    };
+    `}
+       ${n("index")} {
+         if (index < uniforms.size) {
+            let a = A[index];
+            let b = B[0];
+
+            let v4a1 = vec4<f32>(
+                unpack2x16float(u32(a.x)),
+                unpack2x16float(u32(a.y))
+            );
+            let v4a2 = vec4<f32>(
+                unpack2x16float(u32(a.z)),
+                unpack2x16float(u32(a.w))
+            );
+
+            let v4res1 = binaryOperation(v4a1, vec4<f32>(b));
+            let v4res2 = binaryOperation(v4a2, vec4<f32>(b));
+
+            let res = vec4<i32>(
+                i32(pack2x16float(v4res1.xy)),
+                i32(pack2x16float(v4res1.zw)),
+                i32(pack2x16float(v4res2.xy)),
+                i32(pack2x16float(v4res2.zw))
+            );
+
+            result[index] = res;
+         }
+       }
+       `;
+	}
+};
+//#endregion
+export { s as BinaryOpProgram, c as BinaryOpScalarProgram, o as BinaryOpType };

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

@@ -0,0 +1,7 @@
+import { default as WebGPUBackendPatch } from '../../../../patches/webgpu_backend';
+export interface DeviceInformation {
+    subgroupsSupported: boolean;
+    subgroupMaxSize: number;
+    variableSubgroups: boolean;
+}
+export default function createDeviceInformation(backend: WebGPUBackendPatch): DeviceInformation;

package/dist/ops/webgpu/utils/deviceInfo.js

@@ -0,0 +1,11 @@
+//#region lib/ops/webgpu/utils/deviceInfo.ts
+function e(e) {
+	let t = e.device.features.has("subgroups");
+	return {
+		subgroupsSupported: t,
+		subgroupMaxSize: e.subgroupMaxSize,
+		variableSubgroups: e.subgroupMinSize !== e.subgroupMaxSize && t
+	};
+}
+//#endregion
+export { e as default };

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

@@ -0,0 +1,43 @@
+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';
+import { DeviceInformation } from './deviceInfo';
+export interface ReduceParams {
+    reductionOp: 'mean' | 'sum';
+    elementwise?: boolean;
+    forceWorkgroupSize?: number;
+}
+export declare function createReduceInfo(inputs: TensorInfo[], axis: number | number[]): backend_util.ReduceInfo;
+export declare class ReduceProgram implements WebGPUProgram {
+    atomic: boolean;
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    variableNames: string[];
+    uniforms: string;
+    inputShape: number[];
+    size: boolean;
+    packed: boolean;
+    outputComponent: number;
+    variableComponents?: number[];
+    elementwise: boolean;
+    subgroups: boolean;
+    subgroupBuiltins: boolean;
+    deviceInfo: DeviceInformation;
+    params: ReduceParams;
+    utilityFunctions?: string;
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, params: ReduceParams, packed: boolean);
+    protected getWriteSnippet(): string;
+    protected getPreprocessSnippet(): string;
+    protected getPostprocessSnippet(): string;
+    protected getReadSnippet(): string;
+    getUserCode(): string;
+}
+export declare function reduce(program: ReduceProgram, inputs: TensorInfo[], backend: WebGPUBackend, extraUniforms?: {
+    type: string;
+    data: number[];
+}[], output?: TensorInfo): TensorInfo;

package/dist/ops/webgpu/utils/reductions.js

@@ -0,0 +1,263 @@
+import { ii as e, kt as t, nc as n, oc as r, wn as i } from "../../../dist-BewPQWjc.js";
+import { reshape16 as a } from "../../reshape16.js";
+import { s as o } from "../../../webgpu_program-WOyIVMlZ.js";
+import { c as s } from "../../../webgpu_util-B_F3SShA.js";
+//#region lib/ops/webgpu/utils/reductions.ts
+function c(e, t, n, r) {
+	return e && !t ? "\n            bestValue = subgroupAdd(bestValue);\n        " : e ? `
+            bestValue = subgroupAdd(bestValue);
+            let lane = localId.x % subgroupSize;
+            if (lane == 0) {
+                bestValues[localId.x / subgroupSize] = bestValue;
+            }
+            workgroupBarrier();
+            let numSubgroups = ${n} / subgroupSize;
+            bestValue = select(${r ? "vec2<f32>(0.0f)" : "0.0f"}, bestValues[lane], lane < numSubgroups);
+            bestValue = subgroupAdd(bestValue);
+        ` : `
+            bestValues[localId.x] = bestValue;
+            workgroupBarrier();
+
+            var reduceSize = min(u32(Length), ${n}u);
+            for (var currentSize = reduceSize / 2u; reduceSize > 1u;
+                currentSize = reduceSize / 2u) {
+                let interval = DIV_CEIL(reduceSize, 2u);
+                if (localId.x < currentSize) {
+                    let candidate = bestValues[localId.x + interval];
+                    bestValue = bestValue + candidate;
+                    bestValues[localId.x] = bestValue;
+                }
+                reduceSize = interval;
+                workgroupBarrier();
+            }
+
+            bestValue = bestValues[0];
+        `;
+}
+function l(e) {
+	let t = `${e.workgroupSizeX}`, n = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
+           `, r = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
+	return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+            fn readInput(index: i32) -> vec2<f32> {
+                ${e.inputReadSnippet ? e.inputReadSnippet : "\n                let packed = u32(x[index]);\n                return unpack2x16float(packed);\n                "}
+            }
+
+           ${n}
+           ${e.utilityFunctions ?? ""}
+    
+           ${o("index")} {
+                let outputIndex = index / ${t};
+                let offset = outputIndex * uniforms.reduceSize;
+                var bestValue = 0.0f;
+                let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
+    
+                for (var k = tid; k < Length;
+                    k = k + ${t}) {
+                    var candidate = readInput(offset + k);
+                    ${e.inputSnippet}
+                    bestValue = bestValue + candidate.x + candidate.y;
+                }
+
+                ${r}
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};
+
+                ${e.reducedSnippet ? e.reducedSnippet : ""}
+
+                for (var k = tid; k < Length;
+                    k = k + ${t}) {
+                    ${e.outputSnippet}
+                }
+           }
+         `;
+}
+function u(e) {
+	let t = `${e.workgroupSizeX}`, n = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<vec2<f32>, ${e.workgroupSizeX}>;
+           `, r = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !0);
+	return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           fn readInput(index: i32) -> vec2<f32> {
+                ${e.inputReadSnippet ? e.inputReadSnippet : "\n                let packed = u32(x[index]);\n                return unpack2x16float(packed);\n                "}
+            }
+
+            ${n}
+            ${e.utilityFunctions ?? ""}
+    
+           ${o("index")} {
+                let outputIndex = index / ${t};
+                let offset1 = outputIndex * 2 * uniforms.reduceSize;
+                let offset2 = offset1 + uniforms.reduceSize;
+                var bestValue = vec2<f32>(0.0f, 0.0f);
+                let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
+    
+                for (var k = tid; k < Length;
+                    k = k + ${t}) {
+                    var candidate = readInput(offset1 + k);
+                    ${e.inputSnippet}
+                    let bv1 = candidate.x + candidate.y;
+
+                    candidate = readInput(offset2 + k);
+                    ${e.inputSnippet}
+                    let bv2 = candidate.x + candidate.y;
+
+                    bestValue = bestValue + vec2<f32>(bv1, bv2);
+                }
+                ${r}
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};
+
+                ${e.reducedSnippet ?? ""}
+                ${e.outputSnippet}
+           }
+         `;
+}
+function d(e) {
+	return e.elementwise ? l(e) : u(e);
+}
+function f(e) {
+	let t = `${e.workgroupSizeX}`, n = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
+           `, r = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
+	return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           fn readInput(index: i32) -> f32 {
+                ${e.inputReadSnippet ? e.inputReadSnippet : "\n                return x[index];\n                "}
+            }
+
+           ${n}
+              ${e.utilityFunctions ?? ""}
+    
+           ${o("index")} {
+                let outputIndex = index / ${e.workgroupSizeX};
+                let offset = outputIndex * uniforms.reduceSize;
+                var bestValue = 0.0f;
+                let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
+    
+                for (var k = tid; k < Length;
+                    k = k + ${e.workgroupSizeX}) {
+                    var candidate = readInput(offset + k);
+                    ${e.inputSnippet}
+                    bestValue = bestValue + candidate;
+                }
+                ${r}
+
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
+
+                ${e.reducedSnippet}
+
+                ${e.elementwise ? `for (var k = tid; k < Length;
+                    k = k + ${t}) {
+                    ${e.outputSnippet}
+                }` : `${e.outputSnippet}`}
+           }
+         `;
+}
+function p(e, t) {
+	let a = e[0], o = n(t, a.shape), [, s] = i(a.shape, o), c = r(s), l = r(a.shape) / c;
+	return {
+		windowSize: c,
+		inSize: c,
+		batchSize: l,
+		outSize: l
+	};
+}
+var m = class {
+	atomic = !1;
+	outputShape;
+	shaderKey = "reduce16";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	variableNames = ["x"];
+	uniforms = "reduceSize : i32,";
+	inputShape;
+	size = !1;
+	packed = !0;
+	outputComponent;
+	variableComponents;
+	elementwise;
+	subgroups = !1;
+	subgroupBuiltins = !1;
+	deviceInfo;
+	params;
+	utilityFunctions;
+	constructor(e, t, n, r) {
+		this.params = n, this.inputShape = [t.batchSize, t.inSize], this.deviceInfo = e, this.packed = r;
+		let i = n.forceWorkgroupSize ? n.forceWorkgroupSize : t.inSize % 64 == 0 ? 64 : 32;
+		e.subgroupsSupported && !n.forceWorkgroupSize ? (this.workgroupSize = [
+			Math.min(i, e.subgroupMaxSize),
+			1,
+			1
+		], this.subgroups = !0, e.variableSubgroups && (this.subgroupBuiltins = !0)) : this.workgroupSize[0] = i, this.outputShape = n.elementwise ? [t.batchSize, t.inSize] : r ? [t.outSize / 2] : [t.outSize], this.dispatchLayout = s(this.outputShape), this.dispatch = [
+			n.elementwise ? t.batchSize : r ? t.batchSize / 2 : t.batchSize,
+			1,
+			1
+		], this.outputComponent = 1, this.variableComponents = [1], this.elementwise = n.elementwise === !0;
+	}
+	getWriteSnippet() {
+		return this.packed ? "result[outputIndex] = i32(pack2x16float(bestValue));" : "result[outputIndex] = bestValue;";
+	}
+	getPreprocessSnippet() {
+		return "";
+	}
+	getPostprocessSnippet() {
+		return "";
+	}
+	getReadSnippet() {
+		return this.packed ? "\n                let packed = u32(x[index]);\n                return unpack2x16float(packed);\n                " : "return f32(x[index]);";
+	}
+	getUserCode() {
+		let e = this.workgroupSize[0];
+		return this.packed ? d({
+			...this.params,
+			workgroupSizeX: e,
+			subgroups: this.subgroups,
+			variableSubgroups: this.deviceInfo.variableSubgroups,
+			inputReadSnippet: this.getReadSnippet(),
+			inputSnippet: this.getPreprocessSnippet(),
+			outputSnippet: this.getWriteSnippet(),
+			reducedSnippet: this.getPostprocessSnippet(),
+			utilityFunctions: this.utilityFunctions
+		}) : f({
+			...this.params,
+			workgroupSizeX: e,
+			subgroups: this.subgroups,
+			variableSubgroups: this.deviceInfo.variableSubgroups,
+			inputReadSnippet: this.getReadSnippet(),
+			inputSnippet: this.getPreprocessSnippet(),
+			outputSnippet: this.getWriteSnippet(),
+			reducedSnippet: this.getPostprocessSnippet(),
+			utilityFunctions: this.utilityFunctions
+		});
+	}
+};
+function h(n, r, i, o, s) {
+	let c = r[0], l = [{
+		type: "int32",
+		data: [n.inputShape[n.inputShape.length - 1]]
+	}, ...o ?? []], u = s;
+	!s && n.atomic && (u = t(n.outputShape, "int32"));
+	let d = i.runWebGPUProgram(n, r, n.packed ? "packedF16" : n.atomic ? "int32" : "float32", l, u), f = e().makeTensorFromTensorInfo(d);
+	if (n.outputShape.length === 1 && n.outputShape[0] <= 2) return f;
+	let p = a(f, n.elementwise ? c.shape : n.packed ? [...c.shape.slice(0, -2), c.shape[c.shape.length - 2] / 2] : [...c.shape.slice(0, -2), c.shape[c.shape.length - 2]]);
+	return f.dispose(), p;
+}
+//#endregion
+export { m as ReduceProgram, p as createReduceInfo, h as reduce };

package/dist/pack16-Ck-spx_F.js

@@ -0,0 +1,39 @@
+import { Fi as e, ii as t, mr as n } from "./dist-BewPQWjc.js";
+import { isPackedTensor as r } from "./utilities/packed.js";
+//#region lib/ops/grads/unpack16.ts
+var i = {
+	kernelName: "Unpack16",
+	inputsToSave: [],
+	outputsToSave: [],
+	gradFunc: (e) => ({ x: () => s(e) })
+};
+e(i);
+//#endregion
+//#region lib/ops/unpack16.ts
+function a(e, n = 1, i = !1) {
+	if (!r(e)) return e;
+	let a = t().runKernel("Unpack16", { x: e }, { scaling: n });
+	return i && e.dispose(), a;
+}
+//#endregion
+//#region lib/ops/grads/pack16.ts
+var o = {
+	kernelName: "Pack16",
+	inputsToSave: [],
+	outputsToSave: [],
+	gradFunc: (e, t, r) => ({ x: () => {
+		let t = a(e);
+		return r.originalShape && r.padding && r.padding > 0 ? n(t, Array(t.shape.length).fill(0), r.originalShape) : t;
+	} })
+};
+e(o);
+//#endregion
+//#region lib/ops/pack16.ts
+function s(e, n = 1, r = 0) {
+	return t().runKernel("Pack16", { x: e }, {
+		scaling: n,
+		padding: r
+	});
+}
+//#endregion
+export { i, o as n, a as r, s as t };

package/dist/patches/webgpu_backend.d.ts

@@ -0,0 +1,18 @@
+import { DataType, TensorInfo } from '@tensorflow/tfjs-core';
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu';
+import { WebGPUProgram } from './webgpu_program';
+type ProgramUniform = {
+    type: string;
+    data: number[];
+}[];
+interface ExtendedAdapterInfo extends GPUAdapterInfo {
+    subgroupMaxSize: number;
+    subgroupMinSize: number;
+}
+export default class WebGPUBackendPatch extends WebGPUBackend {
+    readonly subgroupMaxSize: number;
+    readonly subgroupMinSize: number;
+    constructor(device: GPUDevice, adapterInfo?: ExtendedAdapterInfo);
+    runWebGPUProgram(program: WebGPUProgram, inputs: TensorInfo[], outputDtype: DataType, programDefinedUniform?: ProgramUniform, output?: TensorInfo): TensorInfo;
+}
+export {};

package/dist/patches/webgpu_backend.js

@@ -0,0 +1,43 @@
+import { Ms as e, Ps as t, Vs as n, oc as r } from "../dist-BewPQWjc.js";
+import { l as i } from "../webgpu_program-WOyIVMlZ.js";
+import { compileProgram as a } from "./webgpu_program.js";
+import { t as o } from "../webgpu-Dt7BMzWz.js";
+//#region lib/patches/webgpu_backend.ts
+var s = (e, n) => {
+	let r = e.limits.maxComputeWorkgroupsPerDimension, i = n.dispatchLayout, a = n.dispatch;
+	if (a.every((e) => e <= r)) return a;
+	t(a[0] > r && i.y === void 0 && i.z === void 0, () => "Dispatch size exceeds WebGPU limits in Y or Z dimension.");
+	let o = Math.ceil(Math.sqrt(a[0]));
+	return o > r ? (o = Math.ceil(Math.cbrt(a[0])), t(o <= r, () => "Total dispatch size exceeds WebGPU maximum."), [
+		o,
+		o,
+		o
+	]) : [
+		o,
+		o,
+		1
+	];
+}, c = class extends o {
+	subgroupMaxSize;
+	subgroupMinSize;
+	constructor(e, t) {
+		super(e, t), this.subgroupMaxSize = t?.subgroupMaxSize ?? 0, this.subgroupMinSize = t?.subgroupMinSize ?? 0;
+	}
+	runWebGPUProgram(t, o, c, l, u) {
+		if (u ||= this.makeTensorInfo(t.outputShape, c), r(u.shape) === 0) return this.tensorMap.get(u.dataId).values = n(u.dtype, 0), u;
+		this.uploadToGPU(u.dataId), t.dispatch = s(this.device, t);
+		let d = o.map((e, n) => {
+			if (e.dtype === "complex64") throw Error("GPGPUProgram does not support complex64 input. For complex64 dtypes, please separate the program into real and imaginary parts.");
+			return this.uploadToGPU(e.dataId), {
+				dtype: this.tensorMap.get(e.dataId).dtype,
+				shape: e.shape,
+				name: t.variableNames[n]
+			};
+		});
+		t.shaderKey = i(t, d, u);
+		let f = e().getBool("WEBGPU_ENGINE_COMPILE_ONLY");
+		return t.shaderKey in this.pipelineCache || (this.pipelineCache[t.shaderKey] = a(this.device, t, d, u, f)), t.pipeline = this.pipelineCache[t.shaderKey], f || this.recordAndSubmit(t, u, o, l), u;
+	}
+};
+//#endregion
+export { c as default };

package/dist/patches/webgpu_base.d.ts

@@ -0,0 +1,21 @@
+/**
+ * @license
+ * Copyright 2022 Google Inc. All Rights Reserved.
+ * 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 interface GPUOptions {
+    powerPreference?: 'low-power' | 'high-performance';
+    disableSubgroups?: boolean;
+}
+export declare function registerWebGPUBackend(options?: GPUOptions): void;

package/dist/patches/webgpu_base.js

@@ -0,0 +1,22 @@
+import { li as e } from "../dist-BewPQWjc.js";
+import t from "./webgpu_backend.js";
+//#region lib/patches/webgpu_base.ts
+function n(n) {
+	e("webgpu", async () => {
+		let e = { powerPreference: n?.powerPreference ?? "high-performance" };
+		console.log("Using custom WebGPU backend with power preference:", e.powerPreference);
+		let r = await navigator.gpu.requestAdapter(e), i = {}, a = [];
+		r.features.has("timestamp-query") && a.push("timestamp-query"), r.features.has("bgra8unorm-storage") && a.push(["bgra8unorm-storage"]), !n?.disableSubgroups && r.features.has("subgroups") && a.push("subgroups"), i.requiredFeatures = a;
+		let o = r.limits;
+		return i.requiredLimits = {
+			maxComputeWorkgroupStorageSize: o.maxComputeWorkgroupStorageSize,
+			maxComputeWorkgroupsPerDimension: o.maxComputeWorkgroupsPerDimension,
+			maxStorageBufferBindingSize: o.maxStorageBufferBindingSize,
+			maxBufferSize: o.maxBufferSize,
+			maxComputeWorkgroupSizeX: o.maxComputeWorkgroupSizeX,
+			maxComputeInvocationsPerWorkgroup: o.maxComputeInvocationsPerWorkgroup
+		}, new t(await r.requestDevice(i), "info" in r ? r.info : "requestAdapterInfo" in r ? await r.requestAdapterInfo() : void 0);
+	}, 3);
+}
+//#endregion
+export { n as registerWebGPUBackend };

package/dist/patches/webgpu_program.d.ts

@@ -0,0 +1,36 @@
+import { DataType, TensorInfo } from '@tensorflow/tfjs-core';
+export declare enum PixelsOpType {
+    FROM_PIXELS = 0,
+    DRAW = 1
+}
+export interface WebGPUProgram {
+    atomic?: boolean;
+    subgroups?: boolean;
+    subgroupBuiltins?: boolean;
+    dispatch: [number, number, number];
+    dispatchLayout: {
+        x: number[];
+        y?: number[];
+        z?: number[];
+    };
+    outputComponent?: number;
+    outputShape: number[];
+    pixelsOpType?: PixelsOpType;
+    shaderKey: string;
+    size?: boolean;
+    uniforms?: string;
+    variableNames: string[];
+    variableComponents?: number[];
+    workgroupSize: [number, number, number];
+    workPerThread?: number;
+    pipeline?: GPUComputePipeline | Promise<GPUComputePipeline>;
+    getUserCode: () => string;
+}
+export declare const compileProgram: (device: GPUDevice, program: WebGPUProgram, inputsData: InputInfo[], output: TensorInfo, parallelCompilation: boolean) => GPUComputePipeline | Promise<GPUComputePipeline>;
+interface InputInfo {
+    dtype: DataType;
+    shape: number[];
+    name: string;
+}
+export declare function getStartHeaderString(useGlobalIndex: boolean, program: WebGPUProgram): string;
+export {};