@jax-js/jax 0.1.2 → 0.1.3

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

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-DeVfWEFS.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) => {
@@ -530,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++) {
@@ -652,4 +660,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 exports.WebGPUBackend = WebGPUBackend;
-//# sourceMappingURL=webgpu-CcGP160M.cjs.map
+//# sourceMappingURL=webgpu-BVns4DbI.cjs.map

package/dist/{webgpu-BGuG58KZ.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-BqymqzuU.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) => {
@@ -530,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++) {
@@ -652,4 +660,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 export { WebGPUBackend };
-//# sourceMappingURL=webgpu-BGuG58KZ.js.map
+//# sourceMappingURL=webgpu-C9iAP5h5.js.map

package/package.json

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