@jax-js/jax 0.0.1 → 0.0.2

package/dist/webgpu-c5Fe8nx8.cjs

@@ -0,0 +1,591 @@
+const require_backend = require('./backend-BK21PBVP.cjs');
+
+//#region src/backend/webgpu.ts
+/** Implementation of `Backend` that uses WebGPU in browsers. */
+var WebGPUBackend = class {
+	type = "webgpu";
+	maxArgs;
+	pipelines;
+	syncReader;
+	buffers;
+	nextSlot;
+	#cachedShaderMap = /* @__PURE__ */ new Map();
+	constructor(device) {
+		this.device = device;
+		if (require_backend.DEBUG >= 3 && device.adapterInfo) console.info("webgpu adapter:", device.adapterInfo.vendor, device.adapterInfo.architecture);
+		this.maxArgs = this.device.limits.maxStorageBuffersPerShaderStage - 1;
+		this.pipelines = new ShaderPipelineCache(device);
+		this.syncReader = new SyncReader(device);
+		this.buffers = /* @__PURE__ */ new Map();
+		this.nextSlot = 1;
+	}
+	malloc(size, initialData) {
+		let buffer;
+		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.unmap();
+			} else {
+				buffer = this.#createBuffer(size);
+				this.device.queue.writeBuffer(buffer, 0, initialData);
+			}
+		} else buffer = this.#createBuffer(size);
+		const slot = this.nextSlot++;
+		this.buffers.set(slot, {
+			buffer,
+			ref: 1
+		});
+		return slot;
+	}
+	incRef(slot) {
+		const buffer = this.buffers.get(slot);
+		if (!buffer) throw new require_backend.SlotError(slot);
+		buffer.ref++;
+	}
+	decRef(slot) {
+		const buffer = this.buffers.get(slot);
+		if (!buffer) throw new require_backend.SlotError(slot);
+		buffer.ref--;
+		if (buffer.ref === 0) {
+			this.buffers.delete(slot);
+			buffer.buffer.destroy();
+		}
+	}
+	async read(slot, start, count) {
+		const buffer = this.#getBuffer(slot);
+		if (start === void 0) start = 0;
+		if (count === void 0) count = buffer.size - start;
+		const staging = this.#createBuffer(count, { read: true });
+		try {
+			const commandEncoder = this.device.createCommandEncoder();
+			commandEncoder.copyBufferToBuffer(buffer, start, staging, 0, count);
+			this.device.queue.submit([commandEncoder.finish()]);
+			await staging.mapAsync(GPUMapMode.READ);
+			const arrayBuffer = staging.getMappedRange();
+			return arrayBuffer.slice();
+		} finally {
+			staging.destroy();
+		}
+	}
+	readSync(slot, start, count) {
+		const buffer = this.#getBuffer(slot);
+		if (start === void 0) start = 0;
+		if (count === void 0) count = buffer.size - start;
+		return this.syncReader.read(buffer, start, count);
+	}
+	#cachedShader(kernel) {
+		const cacheKey = require_backend.FpHash.hash(kernel);
+		let result = this.#cachedShaderMap.get(cacheKey);
+		if (!result) {
+			result = pipelineSource(this.device, kernel);
+			this.#cachedShaderMap.set(cacheKey, result);
+		}
+		return result;
+	}
+	async prepare(kernel) {
+		const { shader, grid } = this.#cachedShader(kernel);
+		const pipeline = await this.pipelines.prepare(shader);
+		return new require_backend.Executable(kernel, {
+			shader,
+			grid,
+			pipeline
+		});
+	}
+	prepareSync(kernel) {
+		const { shader, grid } = this.#cachedShader(kernel);
+		const pipeline = this.pipelines.prepareSync(shader);
+		return new require_backend.Executable(kernel, {
+			shader,
+			grid,
+			pipeline
+		});
+	}
+	dispatch(exe, inputs, outputs) {
+		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
+		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
+		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;
+	}
+	/**
+	* Create a GPU buffer.
+	*
+	* By default, this creates a general-purpose buffer with the given size.
+	*
+	* - If `mapped` is true, initialize the buffer in mapped mode so that it can
+	*   be populated with data from the CPU. (Call `.unmap()` later.)
+	* - If `read` is true, create a staging buffer for returning data to CPU.
+	*   (Call `.mapAsync()` later.)
+	*/
+	#createBuffer(size, { mapped = false, read = false } = {}) {
+		if (read && mapped) throw new Error("mapped and read cannot both be true");
+		const buffer = this.device.createBuffer({
+			size,
+			usage: read ? GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST : GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
+			mappedAtCreation: mapped
+		});
+		return buffer;
+	}
+};
+function dtypeToWgsl(dtype, storage = false) {
+	switch (dtype) {
+		case require_backend.DType.Bool: return storage ? "i32" : "bool";
+		case require_backend.DType.Int32: return "i32";
+		case require_backend.DType.Uint32: return "u32";
+		case require_backend.DType.Float32: return "f32";
+		default: throw new Error(`Unsupported dtype: ${dtype}`);
+	}
+}
+function constToWgsl(dtype, value) {
+	if (dtype === require_backend.DType.Bool) return value ? "true" : "false";
+	if (dtype === require_backend.DType.Int32) return value.toString();
+	if (dtype === require_backend.DType.Uint32) return value.toString() + "u";
+	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;
+	}
+	throw new Error(`Unsupported const dtype: ${dtype}`);
+}
+/**
+* Compiles an expression into WebGPU shader source code.
+*
+* Returns the shader source and the number of workgroups to dispatch along x
+* and y axes, to run the kernel.
+*/
+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 args = Array.from({ length: nargs }, (_, i) => `in${i}`);
+	const shader = [];
+	let indent = "";
+	const pushIndent = Symbol("pushIndent");
+	const popIndent = Symbol("popIndent");
+	const emit = (...lines) => {
+		for (const line of lines) if (line === pushIndent) indent += "  ";
+		else if (line === popIndent) indent = indent.slice(0, -2);
+		else shader.push(line ? indent + line : line);
+	};
+	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);
+	emit("");
+	const usedArgs = Array.from({ length: nargs }, () => null);
+	tune.exp.fold((exp) => {
+		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg] = exp.dtype;
+	});
+	for (let i = 0; i < nargs; i++) {
+		const ty = dtypeToWgsl(usedArgs[i] ?? require_backend.DType.Float32, true);
+		emit(`@group(0) @binding(${i}) var<storage, read> ${args[i]} : array<${ty}>;`);
+	}
+	const resultTy = dtypeToWgsl(kernel.dtype, true);
+	emit(`@group(0) @binding(${nargs}) var<storage, read_write> result : array<${resultTy}>;`);
+	const workgroupSize = require_backend.findPow2(tune.threadCount, 256);
+	const gridSize = Math.ceil(tune.threadCount / workgroupSize);
+	let gridX = gridSize;
+	let gridY = 1;
+	if (gridSize > device.limits.maxComputeWorkgroupsPerDimension) {
+		gridX = 16384;
+		gridY = Math.ceil(gridSize / gridX);
+	}
+	emit("", `@compute @workgroup_size(${workgroupSize})`, "fn main(@builtin(global_invocation_id) id : vec3<u32>) {", pushIndent);
+	if (gridY === 1) emit(`if (id.x >= ${tune.threadCount}) { return; }`, "let gidx: i32 = i32(id.x);");
+	else {
+		const sizeX = gridX * workgroupSize;
+		emit(`if (${sizeX} * id.y + id.x >= ${tune.threadCount}) { return; }`, `let gidx: i32 = i32(${sizeX} * id.y + id.x);`);
+	}
+	let gensymCount = 0;
+	const gensym = () => `alu${gensymCount++}`;
+	for (let i = 0; i < args.length; i++) if (!usedArgs[i]) emit(`_ = &${args[i]};`);
+	const references = /* @__PURE__ */ new Map();
+	const seen = /* @__PURE__ */ new Set();
+	const countReferences = (exp) => {
+		references.set(exp, (references.get(exp) ?? 0) + 1);
+		if (!seen.has(exp)) {
+			seen.add(exp);
+			for (const src of exp.src) countReferences(src);
+		}
+	};
+	const expContext = /* @__PURE__ */ new Map();
+	const gen = (exp) => {
+		if (expContext.has(exp)) return expContext.get(exp);
+		const { op, src, dtype, arg } = exp;
+		let source = "";
+		if (require_backend.AluGroup.Binary.has(op) || require_backend.AluGroup.Compare.has(op)) {
+			const a = gen(src[0]);
+			const b = gen(src[1]);
+			if (op === require_backend.AluOp.Add) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
+			else source = `(${a} + ${b})`;
+			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.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)) {
+			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.Exp) source = `exp(${a})`;
+			else if (op === require_backend.AluOp.Log) source = `log(${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.Threefry2x32) {
+			const x = gensym();
+			const [k0, k1, c0, c1] = src.map((x$1) => require_backend.strip1(gen(x$1)));
+			emit(`let ${x} = threefry2x32(vec2(${k0}, ${k1}), vec2(${c0}, ${c1}));`);
+			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 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]))}]`;
+			if (dtype === require_backend.DType.Bool) source = `(${source} != 0)`;
+		}
+		if (!source) throw new Error(`Missing impl for op: ${op}`);
+		const typeName = dtypeToWgsl(dtype);
+		if ((references.get(exp) ?? 0) > 1) {
+			const name = gensym();
+			expContext.set(exp, name);
+			emit(`let ${name}: ${typeName} = ${require_backend.strip1(source)};`);
+			return name;
+		} else {
+			expContext.set(exp, source);
+			return source;
+		}
+	};
+	if (!kernel.reduction) {
+		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;
+		const acc = [...Array(upcast)].map((_, i) => `acc${i}`);
+		for (let i = 0; i < upcast; i++) emit(`var ${acc[i]}: ${dtypeToWgsl(re.dtype)} = ${constToWgsl(re.dtype, re.identity)};`);
+		emit(`for (var ridx: i32 = 0; ridx < ${tune.size.reduce}; ridx++) {`, pushIndent);
+		const exps = [];
+		const cache = /* @__PURE__ */ new Map();
+		for (let up = 0; up < upcast; up++) {
+			exps.push([]);
+			for (let un = 0; un < unroll; un++) {
+				const exp = tune.exp.substitute({
+					upcast: require_backend.AluExp.i32(up),
+					unroll: require_backend.AluExp.i32(un)
+				});
+				exps[up].push(exp.simplify(cache));
+				countReferences(exps[up][un]);
+			}
+		}
+		const items = exps.map((ar) => ar.map(gen).map(require_backend.strip1));
+		for (let i = 0; i < upcast; i++) {
+			let rhs = items[i][0];
+			for (let j = 1; j < unroll; j++) if (re.op === require_backend.AluOp.Add) rhs = `${rhs} + ${items[i][j]}`;
+			else if (re.op === require_backend.AluOp.Mul) rhs = `${rhs} * ${items[i][j]}`;
+			else if (re.op === require_backend.AluOp.Min) rhs = `min(${rhs}, ${items[i][j]})`;
+			else if (re.op === require_backend.AluOp.Max) rhs = `max(${rhs}, ${items[i][j]})`;
+			else throw new Error(`Unsupported reduction op: ${re.op}`);
+			if (re.op === require_backend.AluOp.Add) emit(`${acc[i]} += ${rhs};`);
+			else if (re.op === require_backend.AluOp.Mul) emit(`${acc[i]} *= ${rhs};`);
+			else if (re.op === require_backend.AluOp.Min) emit(`${acc[i]} = min(${acc[i]}, ${rhs});`);
+			else if (re.op === require_backend.AluOp.Max) emit(`${acc[i]} = max(${acc[i]}, ${rhs});`);
+			else throw new Error(`Unsupported reduction op: ${re.op}`);
+		}
+		emit(popIndent, "}");
+		expContext.clear();
+		references.clear();
+		seen.clear();
+		const outputIdxExps = [];
+		const fusionExps = [];
+		for (let i = 0; i < upcast; i++) {
+			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));
+			countReferences(fusionExps[i]);
+		}
+		for (let i = 0; i < upcast; i++) {
+			const index = require_backend.strip1(gen(outputIdxExps[i]));
+			let rhs = require_backend.strip1(gen(fusionExps[i]));
+			if (resultTy !== dtypeToWgsl(fusionExps[i].dtype)) rhs = `${resultTy}(${rhs})`;
+			emit(`result[${index}] = ${rhs};`);
+		}
+	}
+	emit(popIndent, "}");
+	return {
+		shader: shader.join("\n"),
+		grid: [gridX, gridY]
+	};
+}
+function pipelineSubmit(device, { pipeline, grid }, inputs, outputs) {
+	if (inputs.length + outputs.length > device.limits.maxStorageBuffersPerShaderStage) {
+		const actual = inputs.length + outputs.length;
+		const max = device.limits.maxStorageBuffersPerShaderStage;
+		throw new Error(`Too many buffers (${actual}) for WebGPU pipeline (max: ${max})`);
+	}
+	const bindGroup = device.createBindGroup({
+		layout: pipeline.getBindGroupLayout(0),
+		entries: [...inputs.map((buffer, i) => {
+			return {
+				binding: i,
+				resource: { buffer }
+			};
+		}), {
+			binding: inputs.length,
+			resource: { buffer: outputs[0] }
+		}]
+	});
+	const commandEncoder = device.createCommandEncoder();
+	const passEncoder = commandEncoder.beginComputePass();
+	passEncoder.setPipeline(pipeline);
+	passEncoder.setBindGroup(0, bindGroup);
+	passEncoder.dispatchWorkgroups(grid[0], grid[1]);
+	passEncoder.end();
+	device.queue.submit([commandEncoder.finish()]);
+}
+/**
+* A cache for compiled GPU compute pipelines, keyed by the shader source.
+*
+* This supports both async compilation (recommended) and a synchronous variant.
+* If the pipeline is not in the cache, it will be compiled and added. For async
+* compilation, only one compilation will be in progress at a time for a given
+* shader source.
+*/
+var ShaderPipelineCache = class {
+	cache;
+	inProgress;
+	constructor(device) {
+		this.device = device;
+		this.cache = /* @__PURE__ */ new Map();
+		this.inProgress = /* @__PURE__ */ new Map();
+	}
+	async prepare(code) {
+		const existingPipeline = this.cache.get(code);
+		if (existingPipeline) return existingPipeline;
+		const existingPromise = this.inProgress.get(code);
+		if (existingPromise) return await existingPromise;
+		if (require_backend.DEBUG >= 2) console.info("=========== WebGPU shader ===========\n" + code);
+		const shaderModule = this.device.createShaderModule({ code });
+		const promise = (async () => {
+			this.device.pushErrorScope("validation");
+			try {
+				const pipeline$1 = await this.device.createComputePipelineAsync({
+					layout: "auto",
+					compute: {
+						module: shaderModule,
+						entryPoint: "main"
+					}
+				});
+				await this.device.popErrorScope();
+				return pipeline$1;
+			} catch (_error) {
+				const scope = await this.device.popErrorScope();
+				const emsg = await compileError(shaderModule, scope, code);
+				throw new Error(emsg);
+			}
+		})();
+		this.inProgress.set(code, promise);
+		const pipeline = await promise;
+		this.cache.set(code, pipeline);
+		return pipeline;
+	}
+	prepareSync(code) {
+		const existingPipeline = this.cache.get(code);
+		if (existingPipeline) return existingPipeline;
+		if (require_backend.DEBUG >= 2) console.info("=========== WebGPU shader ===========\n" + code);
+		const shaderModule = this.device.createShaderModule({ code });
+		this.device.pushErrorScope("validation");
+		const pipeline = this.device.createComputePipeline({
+			layout: "auto",
+			compute: {
+				module: shaderModule,
+				entryPoint: "main"
+			}
+		});
+		this.device.popErrorScope().then(async (scope) => {
+			if (scope !== null) {
+				const emsg = await compileError(shaderModule, scope, code);
+				console.error(emsg);
+			}
+		});
+		this.cache.set(code, pipeline);
+		return pipeline;
+	}
+};
+/** Gather information about a compilation error and format it. */
+async function compileError(shaderModule, scope, code) {
+	let message = `Failed to compile shader: ${scope ? scope.message : "(no error scope)"}`;
+	const info = await shaderModule.getCompilationInfo();
+	for (const msg of info.messages) message += `\n  [${msg.type} at ${msg.lineNum}:${msg.linePos}] ${msg.message}`;
+	if (code) message += `\n\n${code}`;
+	return message;
+}
+/**
+* Graphics state used to synchronously read data from WebGPU buffers.
+*
+* This trick is borrowed from TensorFlow.js. Basically, the idea is to create
+* an offscreen canvas with one pixel for every 4 bytes ("device storage"), then
+* configure it with a WebGPU context. Copy the buffer to a texture, then draw
+* the canvas onto another offscreen canvas with '2d' context ("host storage").
+*
+* Once it's on host storage, we can use `getImageData()` to read the pixels
+* from the image directly.
+*
+* We use 256x256 canvases here (256 KiB). The performance of this is bad
+* because it involves multiple data copies, but it still works. We also
+* actually need to copy the image twice: once in "opaque" mode for the RGB
+* values, and once in "premultiplied" mode for the alpha channel.
+*
+* https://github.com/tensorflow/tfjs/blob/tfjs-v4.22.0/tfjs-backend-webgpu/src/backend_webgpu.ts#L379
+*/
+var SyncReader = class SyncReader {
+	static alphaModes = ["opaque", "premultiplied"];
+	static width = 256;
+	static height = 256;
+	initialized = false;
+	deviceStorage;
+	deviceContexts;
+	hostStorage;
+	hostContext;
+	constructor(device) {
+		this.device = device;
+	}
+	#init() {
+		const makeCanvas = () => new OffscreenCanvas(SyncReader.width, SyncReader.height);
+		this.deviceStorage = SyncReader.alphaModes.map(makeCanvas);
+		this.deviceContexts = this.deviceStorage.map((canvas, i) => {
+			const context = canvas.getContext("webgpu");
+			context.configure({
+				device: this.device,
+				format: "bgra8unorm",
+				usage: GPUTextureUsage.COPY_DST,
+				alphaMode: SyncReader.alphaModes[i]
+			});
+			return context;
+		});
+		this.hostStorage = makeCanvas();
+		this.hostContext = this.hostStorage.getContext("2d", { willReadFrequently: true });
+		this.initialized = true;
+	}
+	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 bytesPerRow = SyncReader.width * 4;
+		const valsGPU = new ArrayBuffer(count);
+		for (let i = 0; i < deviceContexts.length; i++) {
+			const texture = deviceContexts[i].getCurrentTexture();
+			const readData = (width, height, offset$1) => {
+				const encoder = this.device.createCommandEncoder();
+				encoder.copyBufferToTexture({
+					buffer,
+					bytesPerRow,
+					offset: offset$1 + start
+				}, { texture }, {
+					width,
+					height,
+					depthOrArrayLayers: 1
+				});
+				const commandBuffer = encoder.finish();
+				this.device.queue.submit([commandBuffer]);
+				hostContext.clearRect(0, 0, width, height);
+				hostContext.drawImage(deviceStorage[i], 0, 0);
+				const values = hostContext.getImageData(0, 0, width, height).data;
+				const span = new Uint8ClampedArray(valsGPU, offset$1, 4 * width * height);
+				const alphaMode = SyncReader.alphaModes[i];
+				for (let k = 0; k < span.length; k += 4) if (alphaMode === "premultiplied") span[k + 3] = values[k + 3];
+				else {
+					span[k] = values[k + 2];
+					span[k + 1] = values[k + 1];
+					span[k + 2] = values[k];
+				}
+			};
+			const pixelsPerCanvas = SyncReader.width * SyncReader.height;
+			const wholeChunks = Math.floor(pixelsSize / pixelsPerCanvas);
+			let remainder = pixelsSize % pixelsPerCanvas;
+			const remainderRows = Math.floor(remainder / SyncReader.width);
+			remainder = remainder % SyncReader.width;
+			let offset = 0;
+			for (let j = 0; j < wholeChunks; j++) {
+				readData(SyncReader.width, SyncReader.height, offset);
+				offset += pixelsPerCanvas * 4;
+			}
+			if (remainderRows > 0) {
+				readData(SyncReader.width, remainderRows, offset);
+				offset += remainderRows * SyncReader.width * 4;
+			}
+			if (remainder > 0) readData(remainder, 1, offset);
+		}
+		return valsGPU;
+	}
+};
+const threefrySrc = `
+fn threefry2x32(key: vec2<u32>, ctr: vec2<u32>) -> vec2<u32> {
+  let ks0: u32 = key.x;
+  let ks1: u32 = key.y;
+  let ks2: u32 = ks0 ^ ks1 ^ 0x1BD11BDAu;
+
+  var x0: u32 = ctr.x + ks0;
+  var x1: u32 = ctr.y + ks1;
+
+  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
+  x0 += ks1;
+  x1 += ks2 + 1u;
+
+  x0 += x1; x1 = (x1 << 17u) | (x1 >> 15u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 29u) | (x1 >> 3u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 16u) | (x1 >> 16u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 24u) | (x1 >> 8u); x1 ^= x0;
+  x0 += ks2;
+  x1 += ks0 + 2u;
+
+  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
+  x0 += ks0;
+  x1 += ks1 + 3u;
+
+  x0 += x1; x1 = (x1 << 17u) | (x1 >> 15u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 29u) | (x1 >> 3u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 16u) | (x1 >> 16u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 24u) | (x1 >> 8u); x1 ^= x0;
+  x0 += ks1;
+  x1 += ks2 + 4u;
+
+  x0 += x1; x1 = (x1 << 13u) | (x1 >> 19u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 15u) | (x1 >> 17u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 26u) | (x1 >> 6u); x1 ^= x0;
+  x0 += x1; x1 = (x1 << 6u) | (x1 >> 26u); x1 ^= x0;
+  x0 += ks2;
+  x1 += ks0 + 5u;
+
+  return vec2<u32>(x0, x1);
+}`;
+
+//#endregion
+exports.WebGPUBackend = WebGPUBackend;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",
@@ -35,35 +35,40 @@
   "types": "dist/index.d.ts",
   "module": "dist/index.js",
   "license": "MIT",
-  "scripts": {
-    "build": "tsup",
-    "build:watch": "tsup --watch src",
-    "format": "prettier --write .",
-    "format:check": "prettier --check .",
-    "lint": "eslint",
-    "test": "vitest",
-    "prepublishOnly": "npm run build",
-    "postpublish": "git tag jax/v$npm_package_version && git push --tags"
-  },
   "devDependencies": {
-    "@eslint/js": "^9.25.0",
+    "@eslint/js": "^9.31.0",
     "@types/debug": "^4.1.12",
-    "@vitest/browser": "^3.0.5",
-    "@webgpu/types": "^0.1.54",
-    "eslint": "^9.25.0",
-    "eslint-plugin-import": "^2.31.0",
+    "@vitest/browser": "^3.2.4",
+    "@webgpu/types": "^0.1.64",
+    "eslint": "^9.31.0",
+    "eslint-plugin-import": "^2.32.0",
     "globals": "^16.0.0",
-    "playwright": "^1.50.1",
-    "prettier": "^3.5.2",
-    "prettier-plugin-svelte": "^3.3.3",
-    "tsup": "^8.3.6",
-    "typescript": "~5.7.3",
-    "typescript-eslint": "^8.30.1",
-    "vitest": "^3.0.5"
+    "playwright": "~1.50.1",
+    "prettier": "^3.6.2",
+    "prettier-plugin-svelte": "^3.4.0",
+    "tsdown": "^0.13.0",
+    "tsx": "^4.20.3",
+    "typedoc": "^0.28.7",
+    "typescript": "~5.8.3",
+    "typescript-eslint": "^8.38.0",
+    "vitest": "^3.2.4"
+  },
+  "engines": {
+    "pnpm": ">=10.0.0"
   },
   "prettier": {
     "plugins": [
       "prettier-plugin-svelte"
     ]
+  },
+  "scripts": {
+    "build": "tsdown",
+    "build:watch": "TSDOWN_WATCH_MODE=1 tsdown",
+    "check": "tsc",
+    "docs": "tsx typedoc.ts",
+    "format": "prettier --write .",
+    "format:check": "prettier --check .",
+    "lint": "eslint",
+    "test": "vitest"
   }
-}
+}