@jax-js/jax 0.1.8 → 0.1.10

package/README.md

@@ -43,6 +43,41 @@ way to get started on a blank HTML page.
 </script>
 ```
 
+### Platforms
+
+This table refers to latest versions of each browser. WebGPU has gained wide support in browsers as
+of late 2025.
+
+| Platform            | CPU (Wasm) | GPU (WebGPU)   | GPU (WebGL) |
+| ------------------- | ---------- | -------------- | ----------- |
+| Chrome / Edge       | ✅         | ✅             | ✅          |
+| Firefox             | ✅         | ✅ - macOS 26+ | ✅          |
+| Safari              | ✅         | ✅ - macOS 26+ | ✅          |
+| iOS                 | ✅         | ✅ - iOS 26+   | ✅          |
+| Chrome for Android  | ✅         | ✅             | ✅          |
+| Firefox for Android | ✅         | ❌             | ✅          |
+| Node.js             | ✅         | ❌             | ❌          |
+| Deno                | ✅         | ✅ - async     | ❌          |
+
+## Examples
+
+Community usage:
+
+- [**autoresearch-webgpu**: autoresesarch, in the browser](https://autoresearch.lucasgelfond.online/)
+- [**tanh.xyz**: Interactive ML visualizations](https://tanh.xyz/)
+- [**jax-js-bayes**: Declarative Bayesian modeling library](https://github.com/StefanSko/jax-js-bayes)
+
+Demos on the jax-js website:
+
+- [Training neural networks on MNIST](https://jax-js.com/mnist)
+- [Voice cloning: Kyutai Pocket TTS](https://jax-js.com/tts)
+- [CLIP embeddings for books in-browser](https://jax-js.com/mobileclip)
+- [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
+- [In-browser REPL](https://jax-js.com/repl)
+- [Matmul benchmark](https://jax-js.com/bench/matmul)
+- [Conv2d benchmark](https://jax-js.com/bench/conv2d)
+- [Mandelbrot set](https://jax-js.com/mandelbrot)
+
 ## Feature comparison
 
 Here's a quick, high-level comparison with other popular web ML runtimes:
@@ -320,6 +355,8 @@ self-contained way in other projects.
 
 ### Performance
 
+To see per-kernel traces in browser development tools, call `jax.profiler.startTrace()`.
+
 The WebGPU runtime includes an ML compiler with tile-aware optimizations, tuned for indiidual
 browsers. Also, this library uniquely has the `jit()` feature that fuses operations together and
 records an execution graph. jax-js achieves **over 7000 GFLOP/s** for matrix multiplication on an
@@ -338,19 +375,6 @@ well as unique optimizations such as FlashAttention variants.
 That's all for this short tutorial. Please see the generated
 [API reference](https://jax-js.com/docs) for detailed documentation.
 
-## Examples
-
-If you make something cool with jax-js, don't be a stranger! We can feature it here.
-
-- [Training neural networks on MNIST](https://jax-js.com/mnist)
-- [Voice cloning: Kyutai Pocket TTS](https://jax-js.com/tts)
-- [CLIP embeddings for books in-browser](https://jax-js.com/mobileclip)
-- [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
-- [In-browser REPL](https://jax-js.com/repl)
-- [Matmul benchmark](https://jax-js.com/bench/matmul)
-- [Conv2d benchmark](https://jax-js.com/bench/conv2d)
-- [Mandelbrot set](https://jax-js.com/mandelbrot)
-
 ## Development
 
 _The following technical details are for contributing to jax-js and modifying its internals._
@@ -363,19 +387,6 @@ pnpm install
 pnpm run build:watch
 ```
 
-The `pnpm install` command automatically sets up Git hooks via
-[Husky](https://typicode.github.io/husky/). Pre-commit hooks will run ESLint and Prettier on staged
-files to ensure code quality.
-
-You can also run linting and formatting manually:
-
-```bash
-pnpm lint          # Run ESLint
-pnpm format        # Format all files with Prettier
-pnpm format:check  # Check formatting without writing
-pnpm check         # Run TypeScript type checking
-```
-
 Then you can run tests in a headless browser using [Vitest](https://vitest.dev/).
 
 ```bash
@@ -392,6 +403,15 @@ To start a Vite dev server running the website, demos and REPL:
 pnpm -C website dev
 ```
 
+You can run the linter, code formatter, and type checker with:
+
+```bash
+pnpm lint          # Run ESLint
+pnpm format        # Format all files with Prettier
+pnpm format:check  # Check formatting without writing
+pnpm check         # Run TypeScript type checking
+```
+
 ## Future work / help wanted
 
 Contributions are welcomed! Some fruitful areas to look into:
@@ -399,9 +419,6 @@ Contributions are welcomed! Some fruitful areas to look into:
 - Adding support for more JAX functions and operations, see [compatibility table](./FEATURES.md).
 - Improving performance of the WebGPU and Wasm runtimes, generating better kernels, and using SIMD
   and multithreading. (Even single-threaded Wasm could be ~20x faster.)
-- Adding support for `jax.profiling`, in particular the start and end trace functions. We should be
-  able to generate `traceEvents` from backends (especially on GPU, with precise timestamp queries)
-  to help with model performance debugging.
 - Helping the JIT compiler to fuse operations in more cases, like `tanh` branches.
 - Making a fast transformer inference engine, comparing against onnxruntime-web.
 

package/dist/{backend-nEolvdLv.js → backend-Ctqs8la1.js}

@@ -1312,6 +1312,10 @@ var Reduction = class {
 		this.epilogue = epilogue;
 		if (!AluGroup.Reduce.has(op)) throw new TypeError(`Unsupported reduction: ${op}`);
 		this.epilogue = epilogue.simplify();
+		if (this.dtype === DType.Float16 && this.op === AluOp.Add) {
+			this.epilogue = this.epilogue.substitute({ acc: AluExp.cast(this.dtype, AluVar.acc(DType.Float32)) });
+			this.dtype = DType.Float32;
+		}
 	}
 	hash(state) {
 		state.update(this.dtype).update(this.op).update(this.size).update(this.epilogue);
@@ -1479,9 +1483,14 @@ var Routine = class {
 };
 /** One of the valid `Routine` that can be dispatched to backend. */
 let Routines = /* @__PURE__ */ function(Routines$1) {
-	/** Stable sorting algorithm along the last axis. */
+	/**
+	* Sort along the last axis.
+	*
+	* This may be _unstable_ but it often doesn't matter, sorting numbers is
+	* bitwise unique up to signed zeros and NaNs.
+	*/
 	Routines$1["Sort"] = "Sort";
-	/** Returns `int32` indices of the stably sorted array. */
+	/** Stable sorting, returns `int32` indices and values of the sorted array. */
 	Routines$1["Argsort"] = "Argsort";
 	/**
 	* Solve a triangular system of equations.
@@ -1545,7 +1554,13 @@ function runArgsort(type, [x], [y, yi]) {
 		const out = y.subarray(offset, offset + n);
 		const outi = yi.subarray(offset, offset + n);
 		for (let i = 0; i < n; i++) outi[i] = i;
-		outi.sort((a, b) => ar[a] - ar[b]);
+		outi.sort((a, b) => {
+			const x$1 = ar[a];
+			const y$1 = ar[b];
+			if (isNaN(x$1)) return isNaN(y$1) ? 0 : 1;
+			if (isNaN(y$1)) return -1;
+			return x$1 === y$1 ? 0 : x$1 < y$1 ? -1 : 1;
+		});
 		for (let i = 0; i < n; i++) out[i] = ar[outi[i]];
 	}
 }
@@ -2255,11 +2270,15 @@ var TuneDims = class {
 };
 /** Tuning step that does not apply any optimization. */
 function tuneNullopt(kernel) {
+	let exp = kernel.exp;
 	const vars = {};
 	vars.gidx = AluExp.special(DType.Int32, "gidx", kernel.size);
-	if (kernel.reduction) vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
+	if (kernel.reduction) {
+		vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
+		if (exp.dtype !== kernel.reduction.dtype) exp = AluExp.cast(kernel.reduction.dtype, exp);
+	}
 	return {
-		exp: kernel.exp.substitute(vars).rewriteGlobalViews().simplify(),
+		exp: exp.substitute(vars).rewriteGlobalViews().simplify(),
 		epilogue: kernel.reduction?.epilogue.substitute({ gidx: vars.gidx }).rewriteGlobalViews().simplify(),
 		outputIdxExp: vars.gidx,
 		threadCount: kernel.size,
@@ -2268,8 +2287,9 @@ function tuneNullopt(kernel) {
 }
 /** Tuning for WebGPU kernels. */
 function tuneWebgpu(kernel) {
-	const { exp, reduction } = kernel;
+	const reduction = kernel.reduction;
 	if (!reduction) return tuneNullopt(kernel);
+	const exp = AluExp.cast(reduction.dtype, kernel.exp);
 	const globalIndexes = exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex);
 	if (globalIndexes.length > 0) {
 		if (DEBUG >= 4) console.info("Tuning: Found GlobalIndex ops, skipping opt.");
@@ -2321,7 +2341,7 @@ function tuneWebgpu(kernel) {
 	if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
 		const s = dim.st.shape[dim.unroll - 1];
 		if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
-		else for (const splits of [8, 4]) if (s % splits === 0) {
+		else for (const splits of [4, 2]) if (s % splits === 0) {
 			dim.applyUnroll(dim.unroll - 1, splits);
 			break;
 		}
@@ -2497,6 +2517,85 @@ var CpuBackend = class {
 	}
 };
 
+//#endregion
+//#region src/tracing.ts
+let traceEnabled = false;
+const flushCallbacks = [];
+/**
+* Start collecting kernel traces.
+*
+* Traces appear in developer tools under the "Performance" tab, and they are
+* useful for measuring fine-grained kernel execution time.
+*/
+function startTrace() {
+	traceEnabled = true;
+}
+/**
+* Stop collecting kernel traces.
+*
+* Traces appear in developer tools under the "Performance" tab, and they are
+* useful for measuring fine-grained kernel execution time.
+*/
+function stopTrace() {
+	traceEnabled = false;
+	for (const cb of flushCallbacks) cb();
+}
+/** Check if tracing is currently enabled. */
+function isTracing() {
+	return traceEnabled;
+}
+/** Register a callback to flush pending trace data when tracing stops. */
+function onFlushTrace(cb) {
+	flushCallbacks.push(cb);
+}
+function humanSize(n) {
+	if (n >= 1e9) return `${(n / 1e9).toPrecision(3)}B`;
+	if (n >= 1e6) return `${(n / 1e6).toPrecision(3)}M`;
+	if (n >= 1e3) return `${(n / 1e3).toPrecision(3)}K`;
+	return `${n}`;
+}
+/** Build a trace label, properties, and color from a kernel or routine source. */
+function traceSourceInfo(source) {
+	const properties = [];
+	let label;
+	let color;
+	if (source instanceof Kernel) {
+		label = `Kernel[${humanSize(source.size)}]`;
+		properties.push(["exp", `${source.exp}`]);
+		properties.push(["size", `${source.size}`]);
+		properties.push(["nargs", `${source.nargs}`]);
+		if (!source.reduction) color = "primary";
+		else {
+			color = "secondary";
+			properties.push(["reduction", `${source.reduction.op}:${source.reduction.size}`]);
+		}
+	} else {
+		color = "tertiary";
+		label = source.name;
+		properties.push(["inputShapes", source.type.inputShapes.map((s) => `[${s}]`).join(", ")]);
+		properties.push(["outputShapes", source.type.outputShapes.map((s) => `[${s}]`).join(", ")]);
+		properties.push(["dtype", source.type.inputDtypes.join(", ")]);
+	}
+	return {
+		label,
+		color,
+		properties
+	};
+}
+/** Emit a trace entry as a `performance.measure` with devtools metadata. */
+function emitTrace(track, info, start, end) {
+	performance.measure(info.label, {
+		detail: { devtools: {
+			trackGroup: "JAX Profiler",
+			track,
+			color: info.color,
+			properties: info.properties
+		} },
+		start,
+		end
+	});
+}
+
 //#endregion
 //#region src/backend/wasm/allocator.ts
 /** Simple tensor memory allocator for WebAssembly linear memory. */
@@ -3903,11 +4002,19 @@ var WasmBackend = class {
 		return new Executable(routine, void 0);
 	}
 	dispatch(exe, inputs, outputs) {
-		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
-		const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
-		const func = instance.exports.kernel;
-		const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
-		func(...ptrs);
+		const tracing = isTracing();
+		const start = tracing ? performance.now() : 0;
+		if (exe.source instanceof Routine) runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
+		else {
+			const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
+			const func = instance.exports.kernel;
+			const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
+			func(...ptrs);
+		}
+		if (tracing) {
+			const info = traceSourceInfo(exe.source);
+			emitTrace("wasm", info, start, performance.now());
+		}
 	}
 	#getBuffer(slot) {
 		const buffer = this.#buffers.get(slot);
@@ -4252,7 +4359,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-B96vzWGE.js");
+		const { WebGPUBackend } = await import("./webgpu-v_W_-oKw.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4290,7 +4397,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await import("./webgl-DweKSWEm.js");
+		const { WebGLBackend } = await import("./webgl-CvQ1QBX1.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4326,4 +4433,4 @@ var UnsupportedRoutineError = class extends Error {
 };
 
 //#endregion
-export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };
+export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, emitTrace, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, isTracing, mapSetUnion, normalizeAxis, onFlushTrace, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, strip1, toposort, traceSourceInfo, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };

package/dist/{backend-B3foXiV_.cjs → backend-DMauYnfl.cjs}

@@ -1313,6 +1313,10 @@ var Reduction = class {
 		this.epilogue = epilogue;
 		if (!AluGroup.Reduce.has(op)) throw new TypeError(`Unsupported reduction: ${op}`);
 		this.epilogue = epilogue.simplify();
+		if (this.dtype === DType.Float16 && this.op === AluOp.Add) {
+			this.epilogue = this.epilogue.substitute({ acc: AluExp.cast(this.dtype, AluVar.acc(DType.Float32)) });
+			this.dtype = DType.Float32;
+		}
 	}
 	hash(state) {
 		state.update(this.dtype).update(this.op).update(this.size).update(this.epilogue);
@@ -1480,9 +1484,14 @@ var Routine = class {
 };
 /** One of the valid `Routine` that can be dispatched to backend. */
 let Routines = /* @__PURE__ */ function(Routines$1) {
-	/** Stable sorting algorithm along the last axis. */
+	/**
+	* Sort along the last axis.
+	*
+	* This may be _unstable_ but it often doesn't matter, sorting numbers is
+	* bitwise unique up to signed zeros and NaNs.
+	*/
 	Routines$1["Sort"] = "Sort";
-	/** Returns `int32` indices of the stably sorted array. */
+	/** Stable sorting, returns `int32` indices and values of the sorted array. */
 	Routines$1["Argsort"] = "Argsort";
 	/**
 	* Solve a triangular system of equations.
@@ -1546,7 +1555,13 @@ function runArgsort(type, [x], [y, yi]) {
 		const out = y.subarray(offset, offset + n);
 		const outi = yi.subarray(offset, offset + n);
 		for (let i = 0; i < n; i++) outi[i] = i;
-		outi.sort((a, b) => ar[a] - ar[b]);
+		outi.sort((a, b) => {
+			const x$1 = ar[a];
+			const y$1 = ar[b];
+			if (isNaN(x$1)) return isNaN(y$1) ? 0 : 1;
+			if (isNaN(y$1)) return -1;
+			return x$1 === y$1 ? 0 : x$1 < y$1 ? -1 : 1;
+		});
 		for (let i = 0; i < n; i++) out[i] = ar[outi[i]];
 	}
 }
@@ -2256,11 +2271,15 @@ var TuneDims = class {
 };
 /** Tuning step that does not apply any optimization. */
 function tuneNullopt(kernel) {
+	let exp = kernel.exp;
 	const vars = {};
 	vars.gidx = AluExp.special(DType.Int32, "gidx", kernel.size);
-	if (kernel.reduction) vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
+	if (kernel.reduction) {
+		vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
+		if (exp.dtype !== kernel.reduction.dtype) exp = AluExp.cast(kernel.reduction.dtype, exp);
+	}
 	return {
-		exp: kernel.exp.substitute(vars).rewriteGlobalViews().simplify(),
+		exp: exp.substitute(vars).rewriteGlobalViews().simplify(),
 		epilogue: kernel.reduction?.epilogue.substitute({ gidx: vars.gidx }).rewriteGlobalViews().simplify(),
 		outputIdxExp: vars.gidx,
 		threadCount: kernel.size,
@@ -2269,8 +2288,9 @@ function tuneNullopt(kernel) {
 }
 /** Tuning for WebGPU kernels. */
 function tuneWebgpu(kernel) {
-	const { exp, reduction } = kernel;
+	const reduction = kernel.reduction;
 	if (!reduction) return tuneNullopt(kernel);
+	const exp = AluExp.cast(reduction.dtype, kernel.exp);
 	const globalIndexes = exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex);
 	if (globalIndexes.length > 0) {
 		if (DEBUG >= 4) console.info("Tuning: Found GlobalIndex ops, skipping opt.");
@@ -2322,7 +2342,7 @@ function tuneWebgpu(kernel) {
 	if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
 		const s = dim.st.shape[dim.unroll - 1];
 		if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
-		else for (const splits of [8, 4]) if (s % splits === 0) {
+		else for (const splits of [4, 2]) if (s % splits === 0) {
 			dim.applyUnroll(dim.unroll - 1, splits);
 			break;
 		}
@@ -2498,6 +2518,85 @@ var CpuBackend = class {
 	}
 };
 
+//#endregion
+//#region src/tracing.ts
+let traceEnabled = false;
+const flushCallbacks = [];
+/**
+* Start collecting kernel traces.
+*
+* Traces appear in developer tools under the "Performance" tab, and they are
+* useful for measuring fine-grained kernel execution time.
+*/
+function startTrace() {
+	traceEnabled = true;
+}
+/**
+* Stop collecting kernel traces.
+*
+* Traces appear in developer tools under the "Performance" tab, and they are
+* useful for measuring fine-grained kernel execution time.
+*/
+function stopTrace() {
+	traceEnabled = false;
+	for (const cb of flushCallbacks) cb();
+}
+/** Check if tracing is currently enabled. */
+function isTracing() {
+	return traceEnabled;
+}
+/** Register a callback to flush pending trace data when tracing stops. */
+function onFlushTrace(cb) {
+	flushCallbacks.push(cb);
+}
+function humanSize(n) {
+	if (n >= 1e9) return `${(n / 1e9).toPrecision(3)}B`;
+	if (n >= 1e6) return `${(n / 1e6).toPrecision(3)}M`;
+	if (n >= 1e3) return `${(n / 1e3).toPrecision(3)}K`;
+	return `${n}`;
+}
+/** Build a trace label, properties, and color from a kernel or routine source. */
+function traceSourceInfo(source) {
+	const properties = [];
+	let label;
+	let color;
+	if (source instanceof Kernel) {
+		label = `Kernel[${humanSize(source.size)}]`;
+		properties.push(["exp", `${source.exp}`]);
+		properties.push(["size", `${source.size}`]);
+		properties.push(["nargs", `${source.nargs}`]);
+		if (!source.reduction) color = "primary";
+		else {
+			color = "secondary";
+			properties.push(["reduction", `${source.reduction.op}:${source.reduction.size}`]);
+		}
+	} else {
+		color = "tertiary";
+		label = source.name;
+		properties.push(["inputShapes", source.type.inputShapes.map((s) => `[${s}]`).join(", ")]);
+		properties.push(["outputShapes", source.type.outputShapes.map((s) => `[${s}]`).join(", ")]);
+		properties.push(["dtype", source.type.inputDtypes.join(", ")]);
+	}
+	return {
+		label,
+		color,
+		properties
+	};
+}
+/** Emit a trace entry as a `performance.measure` with devtools metadata. */
+function emitTrace(track, info, start, end) {
+	performance.measure(info.label, {
+		detail: { devtools: {
+			trackGroup: "JAX Profiler",
+			track,
+			color: info.color,
+			properties: info.properties
+		} },
+		start,
+		end
+	});
+}
+
 //#endregion
 //#region src/backend/wasm/allocator.ts
 /** Simple tensor memory allocator for WebAssembly linear memory. */
@@ -3904,11 +4003,19 @@ var WasmBackend = class {
 		return new Executable(routine, void 0);
 	}
 	dispatch(exe, inputs, outputs) {
-		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
-		const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
-		const func = instance.exports.kernel;
-		const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
-		func(...ptrs);
+		const tracing = isTracing();
+		const start = tracing ? performance.now() : 0;
+		if (exe.source instanceof Routine) runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
+		else {
+			const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
+			const func = instance.exports.kernel;
+			const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
+			func(...ptrs);
+		}
+		if (tracing) {
+			const info = traceSourceInfo(exe.source);
+			emitTrace("wasm", info, start, performance.now());
+		}
 	}
 	#getBuffer(slot) {
 		const buffer = this.#buffers.get(slot);
@@ -4253,7 +4360,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-BykvF26B.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-DMSx7a6M.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4291,7 +4398,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-DIIbKJ0G.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-kvVt7-T7.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4495,6 +4602,12 @@ Object.defineProperty(exports, 'dtypedJsArray', {
     return dtypedJsArray;
   }
 });
+Object.defineProperty(exports, 'emitTrace', {
+  enumerable: true,
+  get: function () {
+    return emitTrace;
+  }
+});
 Object.defineProperty(exports, 'findPow2', {
   enumerable: true,
   get: function () {
@@ -4543,6 +4656,12 @@ Object.defineProperty(exports, 'isPermutation', {
     return isPermutation;
   }
 });
+Object.defineProperty(exports, 'isTracing', {
+  enumerable: true,
+  get: function () {
+    return isTracing;
+  }
+});
 Object.defineProperty(exports, 'mapSetUnion', {
   enumerable: true,
   get: function () {
@@ -4555,6 +4674,12 @@ Object.defineProperty(exports, 'normalizeAxis', {
     return normalizeAxis;
   }
 });
+Object.defineProperty(exports, 'onFlushTrace', {
+  enumerable: true,
+  get: function () {
+    return onFlushTrace;
+  }
+});
 Object.defineProperty(exports, 'partitionList', {
   enumerable: true,
   get: function () {
@@ -4603,6 +4728,18 @@ Object.defineProperty(exports, 'setDebug', {
     return setDebug;
   }
 });
+Object.defineProperty(exports, 'startTrace', {
+  enumerable: true,
+  get: function () {
+    return startTrace;
+  }
+});
+Object.defineProperty(exports, 'stopTrace', {
+  enumerable: true,
+  get: function () {
+    return stopTrace;
+  }
+});
 Object.defineProperty(exports, 'strip1', {
   enumerable: true,
   get: function () {
@@ -4615,6 +4752,12 @@ Object.defineProperty(exports, 'toposort', {
     return toposort;
   }
 });
+Object.defineProperty(exports, 'traceSourceInfo', {
+  enumerable: true,
+  get: function () {
+    return traceSourceInfo;
+  }
+});
 Object.defineProperty(exports, 'tuneNullopt', {
   enumerable: true,
   get: function () {