@jax-js/jax 0.1.9 → 0.1.10

package/README.md

@@ -43,13 +43,31 @@ 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
 
-Cool things that the community has made with jax-js:
+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)
 
-And some more demos on the official website.
+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)
@@ -337,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
@@ -367,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
@@ -396,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:
@@ -403,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-BId79r5b.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);
@@ -2266,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,
@@ -2279,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.");
@@ -2508,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. */
@@ -3914,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);
@@ -4263,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-AN0cG_nB.js");
+		const { WebGPUBackend } = await import("./webgpu-v_W_-oKw.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4301,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-DnGrclTz.js");
+		const { WebGLBackend } = await import("./webgl-CvQ1QBX1.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4337,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-DpI0riom.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);
@@ -2267,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,
@@ -2280,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.");
@@ -2509,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. */
@@ -3915,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);
@@ -4264,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-CdjiJSa7.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-DMSx7a6M.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4302,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-C5NjXc1p.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-kvVt7-T7.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4506,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 () {
@@ -4554,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 () {
@@ -4566,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 () {
@@ -4614,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 () {
@@ -4626,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 () {