@jax-js/jax 0.1.1 → 0.1.2

package/dist/{backend-CoVtc9dx.js → backend-BqymqzuU.js}

@@ -289,10 +289,11 @@ var FpHash = class FpHash {
 };
 /** Run a function while caching it inline inside a `Map`. */
 function runWithCache(cache, key, thunk) {
-	if (cache.has(key)) return cache.get(key);
+	const keyStr = JSON.stringify(key);
+	if (cache.has(keyStr)) return cache.get(keyStr);
 	else {
 		const value = thunk();
-		cache.set(key, value);
+		cache.set(keyStr, value);
 		return value;
 	}
 }
@@ -449,6 +450,14 @@ var AluExp = class AluExp {
 	static sqrt(a) {
 		return new AluExp(AluOp.Sqrt, a.dtype, [a]);
 	}
+	static floor(a) {
+		if (!isFloatDtype(a.dtype)) return a;
+		return new AluExp(AluOp.Floor, a.dtype, [a]);
+	}
+	static ceil(a) {
+		if (!isFloatDtype(a.dtype)) return a;
+		return new AluExp(AluOp.Ceil, a.dtype, [a]);
+	}
 	static reciprocal(a) {
 		return new AluExp(AluOp.Reciprocal, a.dtype, [a]);
 	}
@@ -629,6 +638,12 @@ var AluExp = class AluExp {
 			case AluOp.Sqrt:
 				ret = [Math.sqrt(src[0].min), Math.sqrt(src[0].max)];
 				break;
+			case AluOp.Floor:
+				ret = [Math.floor(src[0].min), Math.floor(src[0].max)];
+				break;
+			case AluOp.Ceil:
+				ret = [Math.ceil(src[0].min), Math.ceil(src[0].max)];
+				break;
 			case AluOp.Reciprocal:
 				if (src[0].min <= 0 && src[0].max >= 0) return [-Infinity, Infinity];
 				ret = [1 / src[0].max, 1 / src[0].min];
@@ -861,7 +876,7 @@ var AluExp = class AluExp {
 			else return p(p(src[0].src[0], src[1]), src[0].src[1]).simplify(cache);
 			if (src[1].op === op && src[1].src[1].op === AluOp.Const) return p(p(src[0], src[1].src[0]), src[1].src[1]).simplify(cache);
 		}
-		if (op === AluOp.Mod || op === AluOp.Idiv && src[1].#isConstInt()) {
+		if ((op === AluOp.Mod || op === AluOp.Idiv) && src[1].#isConstInt()) {
 			const [x, y] = src;
 			{
 				const factors = [];
@@ -951,6 +966,8 @@ var AluExp = class AluExp {
 				case AluOp.Erf: return erf(x);
 				case AluOp.Erfc: return erfc(x);
 				case AluOp.Sqrt: return Math.sqrt(x);
+				case AluOp.Floor: return Math.floor(x);
+				case AluOp.Ceil: return Math.ceil(x);
 				case AluOp.Reciprocal: return 1 / x;
 				case AluOp.Cast: {
 					const wasFloat = isFloatDtype(this.src[0].dtype);
@@ -1140,6 +1157,8 @@ let AluOp = /* @__PURE__ */ function(AluOp$1) {
 	AluOp$1["Erf"] = "Erf";
 	AluOp$1["Erfc"] = "Erfc";
 	AluOp$1["Sqrt"] = "Sqrt";
+	AluOp$1["Floor"] = "Floor";
+	AluOp$1["Ceil"] = "Ceil";
 	AluOp$1["Reciprocal"] = "Reciprocal";
 	AluOp$1["Cast"] = "Cast";
 	AluOp$1["Bitcast"] = "Bitcast";
@@ -1174,6 +1193,8 @@ const AluGroup = {
 		AluOp.Erf,
 		AluOp.Erfc,
 		AluOp.Sqrt,
+		AluOp.Floor,
+		AluOp.Ceil,
 		AluOp.Reciprocal,
 		AluOp.Cast,
 		AluOp.Bitcast
@@ -1201,7 +1222,9 @@ const AluGroup = {
 		AluOp.Erf,
 		AluOp.Erfc,
 		AluOp.Sqrt,
-		AluOp.Reciprocal
+		AluOp.Reciprocal,
+		AluOp.Floor,
+		AluOp.Ceil
 	])
 };
 /** Common variables that can be substituted in expressions. */
@@ -3831,7 +3854,9 @@ function translateExp(cg, funcs, exp, ctx) {
 			else if (op === AluOp.Reciprocal) {
 				const dt = dtyF(cg, op, dtype);
 				dt.const(1), gen(src[0]), dt.div();
-			} else if (op === AluOp.Cast) {
+			} else if (op === AluOp.Floor) gen(src[0]), dtyF(cg, op, dtype).floor();
+			else if (op === AluOp.Ceil) gen(src[0]), dtyF(cg, op, dtype).ceil();
+			else if (op === AluOp.Cast) {
 				gen(src[0]);
 				const dtype0 = src[0].dtype;
 				const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
@@ -3977,7 +4002,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-B3UVme6n.js");
+		const { WebGPUBackend } = await import("./webgpu-BGuG58KZ.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4031,4 +4056,4 @@ var UnsupportedOpError = class extends Error {
 
 //#endregion
 export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, ShapeTracker, SlotError, UnsupportedOpError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, 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, tuneWebgpu, unravelAlu, unzip2, zip, zipn };
-//# sourceMappingURL=backend-CoVtc9dx.js.map
+//# sourceMappingURL=backend-BqymqzuU.js.map

package/dist/{backend-BbrKEB18.cjs → backend-DeVfWEFS.cjs}

@@ -290,10 +290,11 @@ var FpHash = class FpHash {
 };
 /** Run a function while caching it inline inside a `Map`. */
 function runWithCache(cache, key, thunk) {
-	if (cache.has(key)) return cache.get(key);
+	const keyStr = JSON.stringify(key);
+	if (cache.has(keyStr)) return cache.get(keyStr);
 	else {
 		const value = thunk();
-		cache.set(key, value);
+		cache.set(keyStr, value);
 		return value;
 	}
 }
@@ -450,6 +451,14 @@ var AluExp = class AluExp {
 	static sqrt(a) {
 		return new AluExp(AluOp.Sqrt, a.dtype, [a]);
 	}
+	static floor(a) {
+		if (!isFloatDtype(a.dtype)) return a;
+		return new AluExp(AluOp.Floor, a.dtype, [a]);
+	}
+	static ceil(a) {
+		if (!isFloatDtype(a.dtype)) return a;
+		return new AluExp(AluOp.Ceil, a.dtype, [a]);
+	}
 	static reciprocal(a) {
 		return new AluExp(AluOp.Reciprocal, a.dtype, [a]);
 	}
@@ -630,6 +639,12 @@ var AluExp = class AluExp {
 			case AluOp.Sqrt:
 				ret = [Math.sqrt(src[0].min), Math.sqrt(src[0].max)];
 				break;
+			case AluOp.Floor:
+				ret = [Math.floor(src[0].min), Math.floor(src[0].max)];
+				break;
+			case AluOp.Ceil:
+				ret = [Math.ceil(src[0].min), Math.ceil(src[0].max)];
+				break;
 			case AluOp.Reciprocal:
 				if (src[0].min <= 0 && src[0].max >= 0) return [-Infinity, Infinity];
 				ret = [1 / src[0].max, 1 / src[0].min];
@@ -862,7 +877,7 @@ var AluExp = class AluExp {
 			else return p(p(src[0].src[0], src[1]), src[0].src[1]).simplify(cache);
 			if (src[1].op === op && src[1].src[1].op === AluOp.Const) return p(p(src[0], src[1].src[0]), src[1].src[1]).simplify(cache);
 		}
-		if (op === AluOp.Mod || op === AluOp.Idiv && src[1].#isConstInt()) {
+		if ((op === AluOp.Mod || op === AluOp.Idiv) && src[1].#isConstInt()) {
 			const [x, y] = src;
 			{
 				const factors = [];
@@ -952,6 +967,8 @@ var AluExp = class AluExp {
 				case AluOp.Erf: return erf(x);
 				case AluOp.Erfc: return erfc(x);
 				case AluOp.Sqrt: return Math.sqrt(x);
+				case AluOp.Floor: return Math.floor(x);
+				case AluOp.Ceil: return Math.ceil(x);
 				case AluOp.Reciprocal: return 1 / x;
 				case AluOp.Cast: {
 					const wasFloat = isFloatDtype(this.src[0].dtype);
@@ -1141,6 +1158,8 @@ let AluOp = /* @__PURE__ */ function(AluOp$1) {
 	AluOp$1["Erf"] = "Erf";
 	AluOp$1["Erfc"] = "Erfc";
 	AluOp$1["Sqrt"] = "Sqrt";
+	AluOp$1["Floor"] = "Floor";
+	AluOp$1["Ceil"] = "Ceil";
 	AluOp$1["Reciprocal"] = "Reciprocal";
 	AluOp$1["Cast"] = "Cast";
 	AluOp$1["Bitcast"] = "Bitcast";
@@ -1175,6 +1194,8 @@ const AluGroup = {
 		AluOp.Erf,
 		AluOp.Erfc,
 		AluOp.Sqrt,
+		AluOp.Floor,
+		AluOp.Ceil,
 		AluOp.Reciprocal,
 		AluOp.Cast,
 		AluOp.Bitcast
@@ -1202,7 +1223,9 @@ const AluGroup = {
 		AluOp.Erf,
 		AluOp.Erfc,
 		AluOp.Sqrt,
-		AluOp.Reciprocal
+		AluOp.Reciprocal,
+		AluOp.Floor,
+		AluOp.Ceil
 	])
 };
 /** Common variables that can be substituted in expressions. */
@@ -3832,7 +3855,9 @@ function translateExp(cg, funcs, exp, ctx) {
 			else if (op === AluOp.Reciprocal) {
 				const dt = dtyF(cg, op, dtype);
 				dt.const(1), gen(src[0]), dt.div();
-			} else if (op === AluOp.Cast) {
+			} else if (op === AluOp.Floor) gen(src[0]), dtyF(cg, op, dtype).floor();
+			else if (op === AluOp.Ceil) gen(src[0]), dtyF(cg, op, dtype).ceil();
+			else if (op === AluOp.Cast) {
 				gen(src[0]);
 				const dtype0 = src[0].dtype;
 				const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
@@ -3978,7 +4003,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-DGYNVHma.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-CcGP160M.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4325,4 +4350,4 @@ Object.defineProperty(exports, 'zipn', {
     return zipn;
   }
 });
-//# sourceMappingURL=backend-BbrKEB18.cjs.map
+//# sourceMappingURL=backend-DeVfWEFS.cjs.map