@jax-js/jax 0.1.1 → 0.1.3

package/dist/{webgpu-DGYNVHma.cjs → webgpu-BVns4DbI.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-BbrKEB18.cjs');
+const require_backend = require('./backend-CmaidnkQ.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -284,6 +284,8 @@ var WebGPUBackend = class {
 		});
 	}
 	dispatch(exe, inputs, outputs) {
+		if (inputs.length !== exe.kernel.nargs) throw new Error(`webgpu: dispatch with ${inputs.length} inputs, expected ${exe.kernel.nargs}`);
+		if (exe.kernel.size === 0) return;
 		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);
@@ -362,12 +364,12 @@ 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 (tune.exp.some((exp) => exp.dtype === require_backend.DType.Float16) || tune.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); }");
-	const distinctOps = require_backend.mapSetUnion(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	const distinctOps = require_backend.mapSetUnion(tune.exp.distinctOps(), tune.epilogue?.distinctOps());
 	if (distinctOps.has(require_backend.AluOp.Threefry2x32)) emit(threefrySrc);
 	if (distinctOps.has(require_backend.AluOp.Erf) || distinctOps.has(require_backend.AluOp.Erfc)) emit(erfSrc);
 	emit("");
@@ -375,6 +377,9 @@ function pipelineSource(device, kernel) {
 	tune.exp.fold((exp) => {
 		if (exp.op === require_backend.AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
 	});
+	tune.epilogue?.fold((exp) => {
+		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);
 		emit(`@group(0) @binding(${i}) var<storage, read> ${args[i]} : array<${ty}>;`);
@@ -398,7 +403,7 @@ function pipelineSource(device, kernel) {
 	let gensymCount = 0;
 	const gensym = () => `alu${gensymCount++}`;
 	const isGensym = (text) => text.match(/^alu[0-9]+$/);
-	for (let i = 0; i < args.length; i++) if (!usedArgs[i]) emit(`_ = &${args[i]};`);
+	if (args.length > 0) emit(args.map((arg) => `_ = &${arg};`).join(" "));
 	const references = /* @__PURE__ */ new Map();
 	const seen = /* @__PURE__ */ new Set();
 	const countReferences = (exp) => {
@@ -436,18 +441,20 @@ function pipelineSource(device, kernel) {
 			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})`;
+			if (op === require_backend.AluOp.Sin) source = `sin(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Cos) source = `cos(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Asin) source = `asin(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Atan) source = `atan(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Exp) source = `exp(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Log) source = `log(${require_backend.strip1(a)})`;
 			else if (op === require_backend.AluOp.Erf || op === require_backend.AluOp.Erfc) {
 				const funcName = op === require_backend.AluOp.Erf ? "erf" : "erfc";
-				if (dtype !== require_backend.DType.Float32) source = `${dtypeToWgsl(dtype)}(${funcName}(f32(${a})))`;
-				else source = `${funcName}(${a})`;
-			} else if (op === require_backend.AluOp.Sqrt) source = `sqrt(${a})`;
+				if (dtype !== require_backend.DType.Float32) source = `${dtypeToWgsl(dtype)}(${funcName}(f32(${require_backend.strip1(a)})))`;
+				else source = `${funcName}(${require_backend.strip1(a)})`;
+			} else if (op === require_backend.AluOp.Sqrt) source = `sqrt(${require_backend.strip1(a)})`;
 			else if (op === require_backend.AluOp.Reciprocal) source = `(1.0 / ${a})`;
+			else if (op === require_backend.AluOp.Floor) source = `floor(${require_backend.strip1(a)})`;
+			else if (op === require_backend.AluOp.Ceil) source = `ceil(${require_backend.strip1(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)})`;
 		}
@@ -528,7 +535,10 @@ 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.epilogue.substitute({ acc: require_backend.AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
+			fusionExps.push(tune.epilogue.substitute({
+				acc: require_backend.AluExp.variable(re.dtype, acc[i]),
+				upcast: require_backend.AluExp.i32(i)
+			}).simplify(cache));
 			countReferences(fusionExps[i]);
 		}
 		for (let i = 0; i < upcast; i++) {
@@ -650,4 +660,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 exports.WebGPUBackend = WebGPUBackend;
-//# sourceMappingURL=webgpu-DGYNVHma.cjs.map
+//# sourceMappingURL=webgpu-BVns4DbI.cjs.map

package/dist/{webgpu-B3UVme6n.js → webgpu-C9iAP5h5.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, UnsupportedOpError, findPow2, isFloatDtype, mapSetUnion, strip1, tuneWebgpu } from "./backend-CoVtc9dx.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, UnsupportedOpError, findPow2, isFloatDtype, mapSetUnion, strip1, tuneWebgpu } from "./backend-BY8wlLEl.js";
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -284,6 +284,8 @@ var WebGPUBackend = class {
 		});
 	}
 	dispatch(exe, inputs, outputs) {
+		if (inputs.length !== exe.kernel.nargs) throw new Error(`webgpu: dispatch with ${inputs.length} inputs, expected ${exe.kernel.nargs}`);
+		if (exe.kernel.size === 0) return;
 		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);
@@ -362,12 +364,12 @@ 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 (tune.exp.some((exp) => exp.dtype === DType.Float16) || tune.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); }");
-	const distinctOps = mapSetUnion(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	const distinctOps = mapSetUnion(tune.exp.distinctOps(), tune.epilogue?.distinctOps());
 	if (distinctOps.has(AluOp.Threefry2x32)) emit(threefrySrc);
 	if (distinctOps.has(AluOp.Erf) || distinctOps.has(AluOp.Erfc)) emit(erfSrc);
 	emit("");
@@ -375,6 +377,9 @@ function pipelineSource(device, kernel) {
 	tune.exp.fold((exp) => {
 		if (exp.op === AluOp.GlobalIndex) usedArgs[exp.arg[0]] = exp.dtype;
 	});
+	tune.epilogue?.fold((exp) => {
+		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);
 		emit(`@group(0) @binding(${i}) var<storage, read> ${args[i]} : array<${ty}>;`);
@@ -398,7 +403,7 @@ function pipelineSource(device, kernel) {
 	let gensymCount = 0;
 	const gensym = () => `alu${gensymCount++}`;
 	const isGensym = (text) => text.match(/^alu[0-9]+$/);
-	for (let i = 0; i < args.length; i++) if (!usedArgs[i]) emit(`_ = &${args[i]};`);
+	if (args.length > 0) emit(args.map((arg) => `_ = &${arg};`).join(" "));
 	const references = /* @__PURE__ */ new Map();
 	const seen = /* @__PURE__ */ new Set();
 	const countReferences = (exp) => {
@@ -436,18 +441,20 @@ function pipelineSource(device, kernel) {
 			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})`;
+			if (op === AluOp.Sin) source = `sin(${strip1(a)})`;
+			else if (op === AluOp.Cos) source = `cos(${strip1(a)})`;
+			else if (op === AluOp.Asin) source = `asin(${strip1(a)})`;
+			else if (op === AluOp.Atan) source = `atan(${strip1(a)})`;
+			else if (op === AluOp.Exp) source = `exp(${strip1(a)})`;
+			else if (op === AluOp.Log) source = `log(${strip1(a)})`;
 			else if (op === AluOp.Erf || op === AluOp.Erfc) {
 				const funcName = op === AluOp.Erf ? "erf" : "erfc";
-				if (dtype !== DType.Float32) source = `${dtypeToWgsl(dtype)}(${funcName}(f32(${a})))`;
-				else source = `${funcName}(${a})`;
-			} else if (op === AluOp.Sqrt) source = `sqrt(${a})`;
+				if (dtype !== DType.Float32) source = `${dtypeToWgsl(dtype)}(${funcName}(f32(${strip1(a)})))`;
+				else source = `${funcName}(${strip1(a)})`;
+			} else if (op === AluOp.Sqrt) source = `sqrt(${strip1(a)})`;
 			else if (op === AluOp.Reciprocal) source = `(1.0 / ${a})`;
+			else if (op === AluOp.Floor) source = `floor(${strip1(a)})`;
+			else if (op === AluOp.Ceil) source = `ceil(${strip1(a)})`;
 			else if (op === AluOp.Cast) source = `${dtypeToWgsl(dtype)}(${strip1(a)})`;
 			else if (op === AluOp.Bitcast) source = `bitcast<${dtypeToWgsl(dtype)}>(${strip1(a)})`;
 		}
@@ -528,7 +535,10 @@ function pipelineSource(device, kernel) {
 			const exp = tune.outputIdxExp.substitute({ upcast: AluExp.i32(i) });
 			outputIdxExps.push(exp.simplify(cache));
 			countReferences(outputIdxExps[i]);
-			fusionExps.push(re.epilogue.substitute({ acc: AluExp.variable(re.dtype, acc[i]) }).simplify(cache));
+			fusionExps.push(tune.epilogue.substitute({
+				acc: AluExp.variable(re.dtype, acc[i]),
+				upcast: AluExp.i32(i)
+			}).simplify(cache));
 			countReferences(fusionExps[i]);
 		}
 		for (let i = 0; i < upcast; i++) {
@@ -650,4 +660,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 export { WebGPUBackend };
-//# sourceMappingURL=webgpu-B3UVme6n.js.map
+//# sourceMappingURL=webgpu-C9iAP5h5.js.map

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.1.1",
+  "version": "0.1.3",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",
@@ -35,19 +35,6 @@
   "types": "dist/index.d.ts",
   "module": "dist/index.js",
   "license": "MIT",
-  "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",
-    "test:coverage": "vitest run --coverage && open coverage/index.html",
-    "prepublishOnly": "pnpm build",
-    "postpublish": "git tag jax/v$npm_package_version && git push --tags"
-  },
   "devDependencies": {
     "@eslint/js": "^9.31.0",
     "@types/debug": "^4.1.12",
@@ -68,15 +55,9 @@
     "typescript-eslint": "^8.46.4",
     "vitest": "^4.0.9"
   },
-  "packageManager": "pnpm@10.22.0",
   "engines": {
     "pnpm": ">=10.0.0"
   },
-  "pnpm": {
-    "overrides": {
-      "@tensorflow/tfjs-core>@webgpu/types": "^0.1.68"
-    }
-  },
   "prettier": {
     "plugins": [
       "prettier-plugin-svelte"
@@ -92,5 +73,16 @@
       }
     ],
     "proseWrap": "always"
+  },
+  "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",
+    "test:coverage": "vitest run --coverage && open coverage/index.html"
   }
-}
+}