@jax-js/jax 0.0.2 → 0.0.4

package/dist/{webgpu-c5Fe8nx8.cjs → webgpu-BVdMaO9T.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-BK21PBVP.cjs');
+const require_backend = require('./backend-Ss1Mev_-.cjs');
 
 //#region src/backend/webgpu.ts
 /** Implementation of `Backend` that uses WebGPU in browsers. */
@@ -21,20 +21,29 @@ var WebGPUBackend = class {
 	}
 	malloc(size, initialData) {
 		let buffer;
+		const paddedSize = Math.ceil(size / 4) * 4;
 		if (initialData) {
 			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
 			if (initialData.byteLength < 4096) {
-				buffer = this.#createBuffer(size, { mapped: true });
-				new Uint8Array(buffer.getMappedRange()).set(new Uint8Array(initialData));
+				buffer = this.#createBuffer(paddedSize, { mapped: true });
+				new Uint8Array(buffer.getMappedRange(), 0, size).set(initialData);
 				buffer.unmap();
 			} else {
-				buffer = this.#createBuffer(size);
-				this.device.queue.writeBuffer(buffer, 0, initialData);
+				buffer = this.#createBuffer(paddedSize);
+				if (initialData.byteLength % 4 === 0) this.device.queue.writeBuffer(buffer, 0, initialData);
+				else {
+					const aligned = initialData.byteLength - initialData.byteLength % 4;
+					this.device.queue.writeBuffer(buffer, 0, initialData, 0, aligned);
+					const remainder = new Uint8Array(4);
+					remainder.set(initialData.subarray(aligned));
+					this.device.queue.writeBuffer(buffer, aligned, remainder);
+				}
 			}
-		} else buffer = this.#createBuffer(size);
+		} else buffer = this.#createBuffer(paddedSize);
 		const slot = this.nextSlot++;
 		this.buffers.set(slot, {
 			buffer,
+			size,
 			ref: 1
 		});
 		return slot;
@@ -54,25 +63,26 @@ var WebGPUBackend = class {
 		}
 	}
 	async read(slot, start, count) {
-		const buffer = this.#getBuffer(slot);
+		const { buffer, size } = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
-		if (count === void 0) count = buffer.size - start;
-		const staging = this.#createBuffer(count, { read: true });
+		if (count === void 0) count = size - start;
+		const paddedSize = Math.ceil(count / 4) * 4;
+		const staging = this.#createBuffer(paddedSize, { read: true });
 		try {
 			const commandEncoder = this.device.createCommandEncoder();
-			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, count);
+			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, paddedSize);
 			this.device.queue.submit([commandEncoder.finish()]);
 			await staging.mapAsync(GPUMapMode.READ);
 			const arrayBuffer = staging.getMappedRange();
-			return arrayBuffer.slice();
+			return new Uint8Array(arrayBuffer.slice(), 0, count);
 		} finally {
 			staging.destroy();
 		}
 	}
 	readSync(slot, start, count) {
-		const buffer = this.#getBuffer(slot);
+		const { buffer, size } = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
-		if (count === void 0) count = buffer.size - start;
+		if (count === void 0) count = size - start;
 		return this.syncReader.read(buffer, start, count);
 	}
 	#cachedShader(kernel) {
@@ -103,14 +113,17 @@ var WebGPUBackend = class {
 		});
 	}
 	dispatch(exe, inputs, outputs) {
-		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
-		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
+		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot).buffer);
+		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot).buffer);
 		pipelineSubmit(this.device, exe.data, inputBuffers, outputBuffers);
 	}
 	#getBuffer(slot) {
 		const buffer = this.buffers.get(slot);
 		if (!buffer) throw new require_backend.SlotError(slot);
-		return buffer.buffer;
+		return {
+			buffer: buffer.buffer,
+			size: buffer.size
+		};
 	}
 	/**
 	* Create a GPU buffer.
@@ -138,6 +151,7 @@ function dtypeToWgsl(dtype, storage = false) {
 		case require_backend.DType.Int32: return "i32";
 		case require_backend.DType.Uint32: return "u32";
 		case require_backend.DType.Float32: return "f32";
+		case require_backend.DType.Float16: return "f16";
 		default: throw new Error(`Unsupported dtype: ${dtype}`);
 	}
 }
@@ -148,9 +162,12 @@ function constToWgsl(dtype, value) {
 	if (dtype === require_backend.DType.Float32) {
 		if (Number.isNaN(value)) return "nan()";
 		if (!Number.isFinite(value)) return value > 0 ? "inf()" : "-inf()";
-		let s = value.toString();
-		if (!s.includes(".")) s += ".0";
-		return s;
+		return "f32(" + value.toString() + ")";
+	}
+	if (dtype === require_backend.DType.Float16) {
+		if (Number.isNaN(value)) return "f16(nan())";
+		if (!Number.isFinite(value)) return value > 0 ? "f16(inf())" : "f16(-inf())";
+		return "f16(" + value.toString() + ")";
 	}
 	throw new Error(`Unsupported const dtype: ${dtype}`);
 }
@@ -163,7 +180,7 @@ function constToWgsl(dtype, value) {
 function pipelineSource(device, kernel) {
 	const tune = require_backend.tuneWebgpu(kernel);
 	if (require_backend.DEBUG >= 3) console.info(`kernel.exp: ${kernel.exp}\ntune.exp: ${tune.exp}`);
-	const { nargs } = kernel;
+	const { nargs, reduction: re } = kernel;
 	const args = Array.from({ length: nargs }, (_, i) => `in${i}`);
 	const shader = [];
 	let indent = "";
@@ -174,12 +191,17 @@ function pipelineSource(device, kernel) {
 		else if (line === popIndent) indent = indent.slice(0, -2);
 		else shader.push(line ? indent + line : line);
 	};
+	if (tune.exp.some((exp) => exp.dtype === require_backend.DType.Float16) || re?.epilogue.some((exp) => exp.dtype === require_backend.DType.Float16)) {
+		if (!device.features.has("shader-f16")) throw new Error("WebGPU device does not support shader-f16 feature");
+		emit("enable f16;");
+	}
 	emit("fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }", "fn inf() -> f32 { let bits = 0x7f800000u; return bitcast<f32>(bits); }");
-	if (tune.exp.collect((exp) => exp.op === require_backend.AluOp.Threefry2x32).length > 0) emit(threefrySrc);
+	const distinctOps = require_backend.union(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	if (distinctOps.has(require_backend.AluOp.Threefry2x32)) emit(threefrySrc);
 	emit("");
 	const usedArgs = Array.from({ length: nargs }, () => null);
 	tune.exp.fold((exp) => {
-		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg] = exp.dtype;
+		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
 	});
 	for (let i = 0; i < nargs; i++) {
 		const ty = dtypeToWgsl(usedArgs[i] ?? require_backend.DType.Float32, true);
@@ -226,22 +248,29 @@ function pipelineSource(device, kernel) {
 			else if (op === require_backend.AluOp.Sub) source = `(${a} - ${b})`;
 			else if (op === require_backend.AluOp.Mul) if (dtype === require_backend.DType.Bool) source = `(${a} && ${b})`;
 			else source = `(${a} * ${b})`;
-			else if (op === require_backend.AluOp.Idiv) source = require_backend.isFloatDtype(dtype) ? `floor(${a} / ${b})` : `(${a} / ${b})`;
+			else if (op === require_backend.AluOp.Idiv) source = require_backend.isFloatDtype(dtype) ? `trunc(${a} / ${b})` : `(${a} / ${b})`;
 			else if (op === require_backend.AluOp.Mod) source = `(${a} % ${b})`;
 			else if (op === require_backend.AluOp.Min) source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
 			else if (op === require_backend.AluOp.Max) source = `max(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
 			else if (op === require_backend.AluOp.Cmplt) source = `(${a} < ${b})`;
 			else if (op === require_backend.AluOp.Cmpne) source = `(${a} != ${b})`;
-		} else if (require_backend.AluGroup.Unary.has(op)) {
+		} else if (require_backend.AluGroup.Unary.has(op)) if (op === require_backend.AluOp.Reciprocal && src[0].op === require_backend.AluOp.Sqrt) {
+			const a = gen(src[0].src[0]);
+			source = `inverseSqrt(${a})`;
+		} else {
 			const a = gen(src[0]);
 			if (op === require_backend.AluOp.Sin) source = `sin(${a})`;
 			else if (op === require_backend.AluOp.Cos) source = `cos(${a})`;
+			else if (op === require_backend.AluOp.Asin) source = `asin(${a})`;
+			else if (op === require_backend.AluOp.Atan) source = `atan(${a})`;
 			else if (op === require_backend.AluOp.Exp) source = `exp(${a})`;
 			else if (op === require_backend.AluOp.Log) source = `log(${a})`;
+			else if (op === require_backend.AluOp.Sqrt) source = `sqrt(${a})`;
 			else if (op === require_backend.AluOp.Reciprocal) source = `(1.0 / ${a})`;
 			else if (op === require_backend.AluOp.Cast) source = `${dtypeToWgsl(dtype)}(${require_backend.strip1(a)})`;
 			else if (op === require_backend.AluOp.Bitcast) source = `bitcast<${dtypeToWgsl(dtype)}>(${require_backend.strip1(a)})`;
-		} else if (op === require_backend.AluOp.Where) source = `select(${require_backend.strip1(gen(src[2]))}, ${require_backend.strip1(gen(src[1]))}, ${require_backend.strip1(gen(src[0]))})`;
+		}
+		else if (op === require_backend.AluOp.Where) source = `select(${require_backend.strip1(gen(src[2]))}, ${require_backend.strip1(gen(src[1]))}, ${require_backend.strip1(gen(src[0]))})`;
 		else if (op === require_backend.AluOp.Threefry2x32) {
 			const x = gensym();
 			const [k0, k1, c0, c1] = src.map((x$1) => require_backend.strip1(gen(x$1)));
@@ -249,15 +278,15 @@ function pipelineSource(device, kernel) {
 			if (arg === "xor") source = `(${x}.x ^ ${x}.y)`;
 			else if (arg === 0) source = `${x}.x`;
 			else if (arg === 1) source = `${x}.y`;
-			else throw new Error("Invalid Threefry2x32 mode: " + arg);
+			else throw new require_backend.UnsupportedOpError(op, dtype, "webgpu", arg);
 		} else if (op === require_backend.AluOp.Const) return constToWgsl(dtype, arg);
 		else if (op === require_backend.AluOp.Special) return arg[0];
 		else if (op === require_backend.AluOp.Variable) return arg;
 		else if (op === require_backend.AluOp.GlobalIndex) {
-			source = `${args[arg]}[${require_backend.strip1(gen(src[0]))}]`;
+			source = `${args[arg[0]]}[${require_backend.strip1(gen(src[0]))}]`;
 			if (dtype === require_backend.DType.Bool) source = `(${source} != 0)`;
 		}
-		if (!source) throw new Error(`Missing impl for op: ${op}`);
+		if (!source) throw new require_backend.UnsupportedOpError(op, dtype, "webgpu", arg);
 		const typeName = dtypeToWgsl(dtype);
 		if ((references.get(exp) ?? 0) > 1) {
 			const name = gensym();
@@ -269,13 +298,12 @@ function pipelineSource(device, kernel) {
 			return source;
 		}
 	};
-	if (!kernel.reduction) {
+	if (!re) {
 		countReferences(tune.exp);
 		let rhs = require_backend.strip1(gen(tune.exp));
 		if (resultTy !== dtypeToWgsl(tune.exp.dtype)) rhs = `${resultTy}(${rhs})`;
 		emit(`result[gidx] = ${rhs};`);
 	} else {
-		const re = kernel.reduction;
 		if ((tune.size.groups ?? 1) > 1) throw new Error("WebGPU backend does not support group optimization yet");
 		const unroll = tune.size.unroll ?? 1;
 		const upcast = tune.size.upcast ?? 1;
@@ -319,7 +347,7 @@ function pipelineSource(device, kernel) {
 			const exp = tune.outputIdxExp.substitute({ upcast: require_backend.AluExp.i32(i) });
 			outputIdxExps.push(exp.simplify(cache));
 			countReferences(outputIdxExps[i]);
-			fusionExps.push(re.fusion.substitute({ acc: require_backend.AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
+			fusionExps.push(re.epilogue.substitute({ acc: require_backend.AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
 			countReferences(fusionExps[i]);
 		}
 		for (let i = 0; i < upcast; i++) {
@@ -487,13 +515,12 @@ var SyncReader = class SyncReader {
 	}
 	read(buffer, start, count) {
 		if (!this.initialized) this.#init();
-		if (count % 4 !== 0) throw new Error("Read size must be a multiple of 4 bytes");
 		const deviceStorage = this.deviceStorage;
 		const deviceContexts = this.deviceContexts;
 		const hostContext = this.hostContext;
-		const pixelsSize = count / 4;
+		const pixelsSize = Math.ceil(count / 4);
 		const bytesPerRow = SyncReader.width * 4;
-		const valsGPU = new ArrayBuffer(count);
+		const valsGPU = /* @__PURE__ */ new ArrayBuffer(pixelsSize * 4);
 		for (let i = 0; i < deviceContexts.length; i++) {
 			const texture = deviceContexts[i].getCurrentTexture();
 			const readData = (width, height, offset$1) => {
@@ -537,7 +564,7 @@ var SyncReader = class SyncReader {
 			}
 			if (remainder > 0) readData(remainder, 1, offset);
 		}
-		return valsGPU;
+		return new Uint8Array(valsGPU, 0, count);
 	}
 };
 const threefrySrc = `

package/dist/{webgpu-JVpVad6g.js → webgpu-ow0Pn_6q.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, findPow2, isFloatDtype, strip1, tuneWebgpu } from "./backend-1eVbAoaV.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, UnsupportedOpError, findPow2, isFloatDtype, strip1, tuneWebgpu, union } from "./backend-EBRGmEYw.js";
 
 //#region src/backend/webgpu.ts
 /** Implementation of `Backend` that uses WebGPU in browsers. */
@@ -21,20 +21,29 @@ var WebGPUBackend = class {
 	}
 	malloc(size, initialData) {
 		let buffer;
+		const paddedSize = Math.ceil(size / 4) * 4;
 		if (initialData) {
 			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
 			if (initialData.byteLength < 4096) {
-				buffer = this.#createBuffer(size, { mapped: true });
-				new Uint8Array(buffer.getMappedRange()).set(new Uint8Array(initialData));
+				buffer = this.#createBuffer(paddedSize, { mapped: true });
+				new Uint8Array(buffer.getMappedRange(), 0, size).set(initialData);
 				buffer.unmap();
 			} else {
-				buffer = this.#createBuffer(size);
-				this.device.queue.writeBuffer(buffer, 0, initialData);
+				buffer = this.#createBuffer(paddedSize);
+				if (initialData.byteLength % 4 === 0) this.device.queue.writeBuffer(buffer, 0, initialData);
+				else {
+					const aligned = initialData.byteLength - initialData.byteLength % 4;
+					this.device.queue.writeBuffer(buffer, 0, initialData, 0, aligned);
+					const remainder = new Uint8Array(4);
+					remainder.set(initialData.subarray(aligned));
+					this.device.queue.writeBuffer(buffer, aligned, remainder);
+				}
 			}
-		} else buffer = this.#createBuffer(size);
+		} else buffer = this.#createBuffer(paddedSize);
 		const slot = this.nextSlot++;
 		this.buffers.set(slot, {
 			buffer,
+			size,
 			ref: 1
 		});
 		return slot;
@@ -54,25 +63,26 @@ var WebGPUBackend = class {
 		}
 	}
 	async read(slot, start, count) {
-		const buffer = this.#getBuffer(slot);
+		const { buffer, size } = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
-		if (count === void 0) count = buffer.size - start;
-		const staging = this.#createBuffer(count, { read: true });
+		if (count === void 0) count = size - start;
+		const paddedSize = Math.ceil(count / 4) * 4;
+		const staging = this.#createBuffer(paddedSize, { read: true });
 		try {
 			const commandEncoder = this.device.createCommandEncoder();
-			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, count);
+			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, paddedSize);
 			this.device.queue.submit([commandEncoder.finish()]);
 			await staging.mapAsync(GPUMapMode.READ);
 			const arrayBuffer = staging.getMappedRange();
-			return arrayBuffer.slice();
+			return new Uint8Array(arrayBuffer.slice(), 0, count);
 		} finally {
 			staging.destroy();
 		}
 	}
 	readSync(slot, start, count) {
-		const buffer = this.#getBuffer(slot);
+		const { buffer, size } = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
-		if (count === void 0) count = buffer.size - start;
+		if (count === void 0) count = size - start;
 		return this.syncReader.read(buffer, start, count);
 	}
 	#cachedShader(kernel) {
@@ -103,14 +113,17 @@ var WebGPUBackend = class {
 		});
 	}
 	dispatch(exe, inputs, outputs) {
-		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
-		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
+		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot).buffer);
+		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot).buffer);
 		pipelineSubmit(this.device, exe.data, inputBuffers, outputBuffers);
 	}
 	#getBuffer(slot) {
 		const buffer = this.buffers.get(slot);
 		if (!buffer) throw new SlotError(slot);
-		return buffer.buffer;
+		return {
+			buffer: buffer.buffer,
+			size: buffer.size
+		};
 	}
 	/**
 	* Create a GPU buffer.
@@ -138,6 +151,7 @@ function dtypeToWgsl(dtype, storage = false) {
 		case DType.Int32: return "i32";
 		case DType.Uint32: return "u32";
 		case DType.Float32: return "f32";
+		case DType.Float16: return "f16";
 		default: throw new Error(`Unsupported dtype: ${dtype}`);
 	}
 }
@@ -148,9 +162,12 @@ function constToWgsl(dtype, value) {
 	if (dtype === DType.Float32) {
 		if (Number.isNaN(value)) return "nan()";
 		if (!Number.isFinite(value)) return value > 0 ? "inf()" : "-inf()";
-		let s = value.toString();
-		if (!s.includes(".")) s += ".0";
-		return s;
+		return "f32(" + value.toString() + ")";
+	}
+	if (dtype === DType.Float16) {
+		if (Number.isNaN(value)) return "f16(nan())";
+		if (!Number.isFinite(value)) return value > 0 ? "f16(inf())" : "f16(-inf())";
+		return "f16(" + value.toString() + ")";
 	}
 	throw new Error(`Unsupported const dtype: ${dtype}`);
 }
@@ -163,7 +180,7 @@ function constToWgsl(dtype, value) {
 function pipelineSource(device, kernel) {
 	const tune = tuneWebgpu(kernel);
 	if (DEBUG >= 3) console.info(`kernel.exp: ${kernel.exp}\ntune.exp: ${tune.exp}`);
-	const { nargs } = kernel;
+	const { nargs, reduction: re } = kernel;
 	const args = Array.from({ length: nargs }, (_, i) => `in${i}`);
 	const shader = [];
 	let indent = "";
@@ -174,12 +191,17 @@ function pipelineSource(device, kernel) {
 		else if (line === popIndent) indent = indent.slice(0, -2);
 		else shader.push(line ? indent + line : line);
 	};
+	if (tune.exp.some((exp) => exp.dtype === DType.Float16) || re?.epilogue.some((exp) => exp.dtype === DType.Float16)) {
+		if (!device.features.has("shader-f16")) throw new Error("WebGPU device does not support shader-f16 feature");
+		emit("enable f16;");
+	}
 	emit("fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }", "fn inf() -> f32 { let bits = 0x7f800000u; return bitcast<f32>(bits); }");
-	if (tune.exp.collect((exp) => exp.op === AluOp.Threefry2x32).length > 0) emit(threefrySrc);
+	const distinctOps = union(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	if (distinctOps.has(AluOp.Threefry2x32)) emit(threefrySrc);
 	emit("");
 	const usedArgs = Array.from({ length: nargs }, () => null);
 	tune.exp.fold((exp) => {
-		if (exp.op === AluOp.GlobalIndex) usedArgs[exp.arg] = exp.dtype;
+		if (exp.op === AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
 	});
 	for (let i = 0; i < nargs; i++) {
 		const ty = dtypeToWgsl(usedArgs[i] ?? DType.Float32, true);
@@ -226,22 +248,29 @@ function pipelineSource(device, kernel) {
 			else if (op === AluOp.Sub) source = `(${a} - ${b})`;
 			else if (op === AluOp.Mul) if (dtype === DType.Bool) source = `(${a} && ${b})`;
 			else source = `(${a} * ${b})`;
-			else if (op === AluOp.Idiv) source = isFloatDtype(dtype) ? `floor(${a} / ${b})` : `(${a} / ${b})`;
+			else if (op === AluOp.Idiv) source = isFloatDtype(dtype) ? `trunc(${a} / ${b})` : `(${a} / ${b})`;
 			else if (op === AluOp.Mod) source = `(${a} % ${b})`;
 			else if (op === AluOp.Min) source = `min(${strip1(a)}, ${strip1(b)})`;
 			else if (op === AluOp.Max) source = `max(${strip1(a)}, ${strip1(b)})`;
 			else if (op === AluOp.Cmplt) source = `(${a} < ${b})`;
 			else if (op === AluOp.Cmpne) source = `(${a} != ${b})`;
-		} else if (AluGroup.Unary.has(op)) {
+		} else if (AluGroup.Unary.has(op)) if (op === AluOp.Reciprocal && src[0].op === AluOp.Sqrt) {
+			const a = gen(src[0].src[0]);
+			source = `inverseSqrt(${a})`;
+		} else {
 			const a = gen(src[0]);
 			if (op === AluOp.Sin) source = `sin(${a})`;
 			else if (op === AluOp.Cos) source = `cos(${a})`;
+			else if (op === AluOp.Asin) source = `asin(${a})`;
+			else if (op === AluOp.Atan) source = `atan(${a})`;
 			else if (op === AluOp.Exp) source = `exp(${a})`;
 			else if (op === AluOp.Log) source = `log(${a})`;
+			else if (op === AluOp.Sqrt) source = `sqrt(${a})`;
 			else if (op === AluOp.Reciprocal) source = `(1.0 / ${a})`;
 			else if (op === AluOp.Cast) source = `${dtypeToWgsl(dtype)}(${strip1(a)})`;
 			else if (op === AluOp.Bitcast) source = `bitcast<${dtypeToWgsl(dtype)}>(${strip1(a)})`;
-		} else if (op === AluOp.Where) source = `select(${strip1(gen(src[2]))}, ${strip1(gen(src[1]))}, ${strip1(gen(src[0]))})`;
+		}
+		else if (op === AluOp.Where) source = `select(${strip1(gen(src[2]))}, ${strip1(gen(src[1]))}, ${strip1(gen(src[0]))})`;
 		else if (op === AluOp.Threefry2x32) {
 			const x = gensym();
 			const [k0, k1, c0, c1] = src.map((x$1) => strip1(gen(x$1)));
@@ -249,15 +278,15 @@ function pipelineSource(device, kernel) {
 			if (arg === "xor") source = `(${x}.x ^ ${x}.y)`;
 			else if (arg === 0) source = `${x}.x`;
 			else if (arg === 1) source = `${x}.y`;
-			else throw new Error("Invalid Threefry2x32 mode: " + arg);
+			else throw new UnsupportedOpError(op, dtype, "webgpu", arg);
 		} else if (op === AluOp.Const) return constToWgsl(dtype, arg);
 		else if (op === AluOp.Special) return arg[0];
 		else if (op === AluOp.Variable) return arg;
 		else if (op === AluOp.GlobalIndex) {
-			source = `${args[arg]}[${strip1(gen(src[0]))}]`;
+			source = `${args[arg[0]]}[${strip1(gen(src[0]))}]`;
 			if (dtype === DType.Bool) source = `(${source} != 0)`;
 		}
-		if (!source) throw new Error(`Missing impl for op: ${op}`);
+		if (!source) throw new UnsupportedOpError(op, dtype, "webgpu", arg);
 		const typeName = dtypeToWgsl(dtype);
 		if ((references.get(exp) ?? 0) > 1) {
 			const name = gensym();
@@ -269,13 +298,12 @@ function pipelineSource(device, kernel) {
 			return source;
 		}
 	};
-	if (!kernel.reduction) {
+	if (!re) {
 		countReferences(tune.exp);
 		let rhs = strip1(gen(tune.exp));
 		if (resultTy !== dtypeToWgsl(tune.exp.dtype)) rhs = `${resultTy}(${rhs})`;
 		emit(`result[gidx] = ${rhs};`);
 	} else {
-		const re = kernel.reduction;
 		if ((tune.size.groups ?? 1) > 1) throw new Error("WebGPU backend does not support group optimization yet");
 		const unroll = tune.size.unroll ?? 1;
 		const upcast = tune.size.upcast ?? 1;
@@ -319,7 +347,7 @@ function pipelineSource(device, kernel) {
 			const exp = tune.outputIdxExp.substitute({ upcast: AluExp.i32(i) });
 			outputIdxExps.push(exp.simplify(cache));
 			countReferences(outputIdxExps[i]);
-			fusionExps.push(re.fusion.substitute({ acc: AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
+			fusionExps.push(re.epilogue.substitute({ acc: AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
 			countReferences(fusionExps[i]);
 		}
 		for (let i = 0; i < upcast; i++) {
@@ -487,13 +515,12 @@ var SyncReader = class SyncReader {
 	}
 	read(buffer, start, count) {
 		if (!this.initialized) this.#init();
-		if (count % 4 !== 0) throw new Error("Read size must be a multiple of 4 bytes");
 		const deviceStorage = this.deviceStorage;
 		const deviceContexts = this.deviceContexts;
 		const hostContext = this.hostContext;
-		const pixelsSize = count / 4;
+		const pixelsSize = Math.ceil(count / 4);
 		const bytesPerRow = SyncReader.width * 4;
-		const valsGPU = new ArrayBuffer(count);
+		const valsGPU = /* @__PURE__ */ new ArrayBuffer(pixelsSize * 4);
 		for (let i = 0; i < deviceContexts.length; i++) {
 			const texture = deviceContexts[i].getCurrentTexture();
 			const readData = (width, height, offset$1) => {
@@ -537,7 +564,7 @@ var SyncReader = class SyncReader {
 			}
 			if (remainder > 0) readData(remainder, 1, offset);
 		}
-		return valsGPU;
+		return new Uint8Array(valsGPU, 0, count);
 	}
 };
 const threefrySrc = `

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.0.2",
+  "version": "0.0.4",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",
@@ -38,20 +38,21 @@
   "devDependencies": {
     "@eslint/js": "^9.31.0",
     "@types/debug": "^4.1.12",
-    "@vitest/browser": "^3.2.4",
+    "@vitest/browser-playwright": "^4.0.9",
     "@webgpu/types": "^0.1.64",
     "eslint": "^9.31.0",
     "eslint-plugin-import": "^2.32.0",
     "globals": "^16.0.0",
-    "playwright": "~1.50.1",
+    "playwright": "~1.52.0",
     "prettier": "^3.6.2",
     "prettier-plugin-svelte": "^3.4.0",
-    "tsdown": "^0.13.0",
+    "tsdown": "^0.13.2",
     "tsx": "^4.20.3",
-    "typedoc": "^0.28.7",
-    "typescript": "~5.8.3",
-    "typescript-eslint": "^8.38.0",
-    "vitest": "^3.2.4"
+    "typedoc": "^0.28.14",
+    "typedoc-theme-fresh": "^0.2.1",
+    "typescript": "~5.9.3",
+    "typescript-eslint": "^8.46.4",
+    "vitest": "^4.0.9"
   },
   "engines": {
     "pnpm": ">=10.0.0"
@@ -59,7 +60,18 @@
   "prettier": {
     "plugins": [
       "prettier-plugin-svelte"
-    ]
+    ],
+    "overrides": [
+      {
+        "files": [
+          "*.md"
+        ],
+        "options": {
+          "printWidth": 100
+        }
+      }
+    ],
+    "proseWrap": "always"
   },
   "scripts": {
     "build": "tsdown",