@jax-js/jax 0.1.12 → 0.1.13

package/README.md

@@ -76,6 +76,7 @@ Demos on the jax-js website:
 - [Object detection: D-FINE (ONNX)](https://jax-js.com/d-fine)
 - [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
 - [Fluid simulation (Navier-Stokes)](https://jax-js.com/fluid-sim)
+- [Neural cellular automata](https://jax-js.com/nca-growing)
 - [In-browser REPL](https://jax-js.com/repl)
 - [Matmul benchmark](https://jax-js.com/bench/matmul)
 - [Matvec benchmark](https://jax-js.com/bench/matvec)

package/dist/{backend-DI-V78Rk.js → backend-DLEk-B3V.js}

@@ -4522,7 +4522,7 @@ var WasmBackend = class {
 		const buffer = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
 		if (count === void 0) count = buffer.byteLength - start;
-		if (buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
+		if (hasSharedArrayBuffer() && buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
 		else return buffer.slice(start, start + count);
 	}
 	async prepareKernel(kernel) {
@@ -5061,7 +5061,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-C2kLdkUh.js");
+		const { WebGPUBackend } = await import("./webgpu-NkF1TZ0t.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -5099,7 +5099,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await import("./webgl-BhsnpeB0.js");
+		const { WebGLBackend } = await import("./webgl-NsFtyIts.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

package/dist/{backend-x-6vqzIM.cjs → backend-DMyuoWi2.cjs}

@@ -4523,7 +4523,7 @@ var WasmBackend = class {
 		const buffer = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
 		if (count === void 0) count = buffer.byteLength - start;
-		if (buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
+		if (hasSharedArrayBuffer() && buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
 		else return buffer.slice(start, start + count);
 	}
 	async prepareKernel(kernel) {
@@ -5062,7 +5062,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-C4S8Uq9e.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-DDGCYtHa.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -5100,7 +5100,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-CD3WK_Me.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-pbfUGDA6.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod$1, isNodeMode, target) => (target = mod$1 != null ? __create(
 }) : target, mod$1));
 
 //#endregion
-const require_backend = require('./backend-x-6vqzIM.cjs');
+const require_backend = require('./backend-DMyuoWi2.cjs');
 
 //#region src/frontend/convolution.ts
 /**

package/dist/index.js

@@ -1,5 +1,5 @@
 import { __export } from "./chunk-Cl8Af3a2.js";
-import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, getWebGPUDevice, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DI-V78Rk.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, getWebGPUDevice, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DLEk-B3V.js";
 
 //#region src/frontend/convolution.ts
 /**

package/dist/{webgl-BhsnpeB0.js → webgl-NsFtyIts.js}

@@ -1,4 +1,4 @@
-import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-DI-V78Rk.js";
+import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-DLEk-B3V.js";
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `

package/dist/{webgl-CD3WK_Me.cjs → webgl-pbfUGDA6.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-x-6vqzIM.cjs');
+const require_backend = require('./backend-DMyuoWi2.cjs');
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `

package/dist/{webgpu-C4S8Uq9e.cjs → webgpu-DDGCYtHa.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-x-6vqzIM.cjs');
+const require_backend = require('./backend-DMyuoWi2.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -1099,6 +1099,8 @@ function flushTracingBatch(device, batch) {
 
 //#endregion
 //#region src/backend/webgpu.ts
+const MAX_REUSABLE_BUFFER_BYTES = 64 * 1024 * 1024;
+const MAX_REUSABLE_BUFFERS_PER_SIZE = 64;
 /** Implementation of `Backend` that uses WebGPU in browsers. */
 var WebGPUBackend = class {
 	type = "webgpu";
@@ -1109,6 +1111,7 @@ var WebGPUBackend = class {
 	nextSlot;
 	#cachedShaderMap = /* @__PURE__ */ new Map();
 	#reusableZsb;
+	#bufferPool = /* @__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);
@@ -1123,31 +1126,22 @@ var WebGPUBackend = class {
 		});
 	}
 	malloc(size, initialData) {
-		let buffer;
-		const paddedSize = Math.ceil(size / 4) * 4;
-		if (size === 0) buffer = this.#reusableZsb;
-		else if (initialData) {
-			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
-			if (initialData.byteLength < 4096) {
-				buffer = this.#createBuffer(paddedSize, { mapped: true });
-				new Uint8Array(buffer.getMappedRange(), 0, size).set(initialData);
-				buffer.unmap();
-			} else {
-				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(paddedSize);
+		if (initialData && initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
+		const allocatedSize = Math.ceil(size / 4) * 4 || 4;
+		const buffer = size === 0 ? this.#reusableZsb : this.#acquireBuffer(allocatedSize);
+		if (initialData && size > 0) if (initialData.byteLength % 4 === 0) this.device.queue.writeBuffer(buffer, 0, initialData);
+		else {
+			const aligned = initialData.byteLength - initialData.byteLength % 4;
+			if (aligned > 0) 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);
+		}
 		const slot = this.nextSlot++;
 		this.buffers.set(slot, {
 			buffer,
 			size,
+			allocatedSize,
 			ref: 1
 		});
 		return slot;
@@ -1163,7 +1157,7 @@ var WebGPUBackend = class {
 		buffer.ref--;
 		if (buffer.ref === 0) {
 			this.buffers.delete(slot);
-			if (buffer.buffer !== this.#reusableZsb) buffer.buffer.destroy();
+			if (buffer.buffer !== this.#reusableZsb) this.#releaseBuffer(buffer.buffer, buffer.allocatedSize);
 		}
 	}
 	async read(slot, start, count) {
@@ -1251,6 +1245,29 @@ var WebGPUBackend = class {
 			size: buffer.size
 		};
 	}
+	#acquireBuffer(size) {
+		if (size > MAX_REUSABLE_BUFFER_BYTES) return this.#createBuffer(size);
+		const bucket = this.#bufferPool.get(size);
+		const buffer = bucket?.pop();
+		if (bucket && bucket.length === 0) this.#bufferPool.delete(size);
+		return buffer ?? this.#createBuffer(size);
+	}
+	#releaseBuffer(buffer, size) {
+		if (size > MAX_REUSABLE_BUFFER_BYTES) {
+			buffer.destroy();
+			return;
+		}
+		const bucket = this.#bufferPool.get(size);
+		if (!bucket) {
+			this.#bufferPool.set(size, [buffer]);
+			return;
+		}
+		if (bucket.length >= MAX_REUSABLE_BUFFERS_PER_SIZE) {
+			buffer.destroy();
+			return;
+		}
+		bucket.push(buffer);
+	}
 	/**
 	* Create a GPU buffer.
 	*

package/dist/{webgpu-C2kLdkUh.js → webgpu-NkF1TZ0t.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, emitTrace, findPow2, isFloatDtype, isTracing, mapSetUnion, onFlushTrace, prod, range, strip1, traceSourceInfo, tuneWebgpu } from "./backend-DI-V78Rk.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, emitTrace, findPow2, isFloatDtype, isTracing, mapSetUnion, onFlushTrace, prod, range, strip1, traceSourceInfo, tuneWebgpu } from "./backend-DLEk-B3V.js";
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -1099,6 +1099,8 @@ function flushTracingBatch(device, batch) {
 
 //#endregion
 //#region src/backend/webgpu.ts
+const MAX_REUSABLE_BUFFER_BYTES = 64 * 1024 * 1024;
+const MAX_REUSABLE_BUFFERS_PER_SIZE = 64;
 /** Implementation of `Backend` that uses WebGPU in browsers. */
 var WebGPUBackend = class {
 	type = "webgpu";
@@ -1109,6 +1111,7 @@ var WebGPUBackend = class {
 	nextSlot;
 	#cachedShaderMap = /* @__PURE__ */ new Map();
 	#reusableZsb;
+	#bufferPool = /* @__PURE__ */ new Map();
 	constructor(device) {
 		this.device = device;
 		if (DEBUG >= 3 && device.adapterInfo) console.info("webgpu adapter:", device.adapterInfo.vendor, device.adapterInfo.architecture);
@@ -1123,31 +1126,22 @@ var WebGPUBackend = class {
 		});
 	}
 	malloc(size, initialData) {
-		let buffer;
-		const paddedSize = Math.ceil(size / 4) * 4;
-		if (size === 0) buffer = this.#reusableZsb;
-		else if (initialData) {
-			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
-			if (initialData.byteLength < 4096) {
-				buffer = this.#createBuffer(paddedSize, { mapped: true });
-				new Uint8Array(buffer.getMappedRange(), 0, size).set(initialData);
-				buffer.unmap();
-			} else {
-				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(paddedSize);
+		if (initialData && initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
+		const allocatedSize = Math.ceil(size / 4) * 4 || 4;
+		const buffer = size === 0 ? this.#reusableZsb : this.#acquireBuffer(allocatedSize);
+		if (initialData && size > 0) if (initialData.byteLength % 4 === 0) this.device.queue.writeBuffer(buffer, 0, initialData);
+		else {
+			const aligned = initialData.byteLength - initialData.byteLength % 4;
+			if (aligned > 0) 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);
+		}
 		const slot = this.nextSlot++;
 		this.buffers.set(slot, {
 			buffer,
 			size,
+			allocatedSize,
 			ref: 1
 		});
 		return slot;
@@ -1163,7 +1157,7 @@ var WebGPUBackend = class {
 		buffer.ref--;
 		if (buffer.ref === 0) {
 			this.buffers.delete(slot);
-			if (buffer.buffer !== this.#reusableZsb) buffer.buffer.destroy();
+			if (buffer.buffer !== this.#reusableZsb) this.#releaseBuffer(buffer.buffer, buffer.allocatedSize);
 		}
 	}
 	async read(slot, start, count) {
@@ -1251,6 +1245,29 @@ var WebGPUBackend = class {
 			size: buffer.size
 		};
 	}
+	#acquireBuffer(size) {
+		if (size > MAX_REUSABLE_BUFFER_BYTES) return this.#createBuffer(size);
+		const bucket = this.#bufferPool.get(size);
+		const buffer = bucket?.pop();
+		if (bucket && bucket.length === 0) this.#bufferPool.delete(size);
+		return buffer ?? this.#createBuffer(size);
+	}
+	#releaseBuffer(buffer, size) {
+		if (size > MAX_REUSABLE_BUFFER_BYTES) {
+			buffer.destroy();
+			return;
+		}
+		const bucket = this.#bufferPool.get(size);
+		if (!bucket) {
+			this.#bufferPool.set(size, [buffer]);
+			return;
+		}
+		if (bucket.length >= MAX_REUSABLE_BUFFERS_PER_SIZE) {
+			buffer.destroy();
+			return;
+		}
+		bucket.push(buffer);
+	}
 	/**
 	* Create a GPU buffer.
 	*

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.1.12",
+  "version": "0.1.13",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",