@jax-js/jax 0.1.0 → 0.1.2

package/dist/{webgpu-LGi2A3mS.js → webgpu-BGuG58KZ.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, UnsupportedOpError, findPow2, isFloatDtype, mapSetUnion, strip1, tuneWebgpu } from "./backend-DwIAd0AG.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, SlotError, UnsupportedOpError, findPow2, isFloatDtype, mapSetUnion, strip1, tuneWebgpu } from "./backend-BqymqzuU.js";
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -323,7 +323,7 @@ function dtypeToWgsl(dtype, storage = false) {
 		case DType.Uint32: return "u32";
 		case DType.Float32: return "f32";
 		case DType.Float16: return "f16";
-		default: throw new Error(`Unsupported dtype: ${dtype}`);
+		default: throw new Error(`Unsupported dtype for WebGPU: ${dtype}`);
 	}
 }
 function constToWgsl(dtype, value) {
@@ -397,6 +397,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]};`);
 	const references = /* @__PURE__ */ new Map();
 	const seen = /* @__PURE__ */ new Set();
@@ -425,24 +426,30 @@ function pipelineSource(device, kernel) {
 			else if (op === AluOp.Min) source = `min(${strip1(a)}, ${strip1(b)})`;
 			else if (op === AluOp.Max) source = `max(${strip1(a)}, ${strip1(b)})`;
 			else if (op === AluOp.Cmplt) source = `(${a} < ${b})`;
-			else if (op === AluOp.Cmpne) source = `(${a} != ${b})`;
+			else if (op === AluOp.Cmpne) if (isFloatDtype(src[0].dtype)) {
+				const x = isGensym(a) ? a : gensym();
+				if (x !== a) emit(`let ${x} = ${a};`);
+				source = `(${x} != ${b} || min(${x}, ${dtypeToWgsl(src[0].dtype)}(inf())) != ${x})`;
+			} else source = `(${a} != ${b})`;
 		} else if (AluGroup.Unary.has(op)) if (op === AluOp.Reciprocal && src[0].op === AluOp.Sqrt) {
 			const a = gen(src[0].src[0]);
 			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)})`;
 		}
@@ -645,4 +652,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 export { WebGPUBackend };
-//# sourceMappingURL=webgpu-LGi2A3mS.js.map
+//# sourceMappingURL=webgpu-BGuG58KZ.js.map

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

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-FtkbO6pI.cjs');
+const require_backend = require('./backend-DeVfWEFS.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -323,7 +323,7 @@ function dtypeToWgsl(dtype, storage = false) {
 		case require_backend.DType.Uint32: return "u32";
 		case require_backend.DType.Float32: return "f32";
 		case require_backend.DType.Float16: return "f16";
-		default: throw new Error(`Unsupported dtype: ${dtype}`);
+		default: throw new Error(`Unsupported dtype for WebGPU: ${dtype}`);
 	}
 }
 function constToWgsl(dtype, value) {
@@ -397,6 +397,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]};`);
 	const references = /* @__PURE__ */ new Map();
 	const seen = /* @__PURE__ */ new Set();
@@ -425,24 +426,30 @@ function pipelineSource(device, kernel) {
 			else if (op === require_backend.AluOp.Min) source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
 			else if (op === require_backend.AluOp.Max) source = `max(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
 			else if (op === require_backend.AluOp.Cmplt) source = `(${a} < ${b})`;
-			else if (op === require_backend.AluOp.Cmpne) source = `(${a} != ${b})`;
+			else if (op === require_backend.AluOp.Cmpne) if (require_backend.isFloatDtype(src[0].dtype)) {
+				const x = isGensym(a) ? a : gensym();
+				if (x !== a) emit(`let ${x} = ${a};`);
+				source = `(${x} != ${b} || min(${x}, ${dtypeToWgsl(src[0].dtype)}(inf())) != ${x})`;
+			} else source = `(${a} != ${b})`;
 		} else if (require_backend.AluGroup.Unary.has(op)) if (op === require_backend.AluOp.Reciprocal && src[0].op === require_backend.AluOp.Sqrt) {
 			const a = gen(src[0].src[0]);
 			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)})`;
 		}
@@ -645,4 +652,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 exports.WebGPUBackend = WebGPUBackend;
-//# sourceMappingURL=webgpu-BE7zA_01.cjs.map
+//# sourceMappingURL=webgpu-CcGP160M.cjs.map

package/package.json

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