@jax-js/jax 0.1.0 → 0.1.1

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

@@ -306,6 +306,7 @@ let DType = /* @__PURE__ */ function(DType$1) {
 	DType$1["Uint32"] = "uint32";
 	DType$1["Bool"] = "bool";
 	DType$1["Float16"] = "float16";
+	DType$1["Float64"] = "float64";
 	return DType$1;
 }({});
 const byteWidth = (dtype) => {
@@ -315,10 +316,11 @@ const byteWidth = (dtype) => {
 		case DType.Uint32:
 		case DType.Bool: return 4;
 		case DType.Float16: return 2;
+		case DType.Float64: return 8;
 		default: throw new TypeError(`Unknown dtype: ${dtype}`);
 	}
 };
-const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float16;
+const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float16 || dtype === DType.Float64;
 /**
 * Promote two dtypes to their join according to the type lattice.
 *
@@ -328,7 +330,7 @@ const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float
 *
 * **Type lattice:**
 * ```text
-* bool -> uint32 -> int32 -> float16 -> float32
+* bool -> uint32 -> int32 -> float16 -> float32 -> float64
 *  weakType --^
 * ```
 *
@@ -350,7 +352,8 @@ function promoteTypes(dtype1, dtype2) {
 		[DType.Uint32]: 1,
 		[DType.Int32]: 2,
 		[DType.Float16]: 3,
-		[DType.Float32]: 4
+		[DType.Float32]: 4,
+		[DType.Float64]: 5
 	};
 	return rank[dtype1] > rank[dtype2] ? dtype1 : dtype2;
 }
@@ -363,6 +366,7 @@ function dtypedArray(dtype, data) {
 		case DType.Bool: return new Int32Array(buffer, byteOffset, length);
 		case DType.Uint32: return new Uint32Array(buffer, byteOffset, length);
 		case DType.Float16: return new Float16Array(buffer, byteOffset, length);
+		case DType.Float64: return new Float64Array(buffer, byteOffset, length);
 		default: throw new Error(`Unimplemented dtype: ${dtype}`);
 	}
 }
@@ -373,6 +377,7 @@ function dtypedJsArray(dtype, data) {
 		case DType.Bool: return new Int32Array(data);
 		case DType.Uint32: return new Uint32Array(data);
 		case DType.Float16: return new Float16Array(data);
+		case DType.Float64: return new Float64Array(data);
 		default: throw new Error(`Unimplemented dtype: ${dtype}`);
 	}
 }
@@ -510,6 +515,9 @@ var AluExp = class AluExp {
 	static f16(value) {
 		return AluExp.const(DType.Float16, value);
 	}
+	static f64(value) {
+		return AluExp.const(DType.Float64, value);
+	}
 	not() {
 		if (this.dtype !== DType.Bool) throw new Error("not() can only be called on boolean expressions");
 		return AluExp.cmpne(this, AluExp.const(DType.Bool, true));
@@ -520,7 +528,8 @@ var AluExp = class AluExp {
 		const hasher = new FpHash();
 		hasher.update(this.op);
 		hasher.update(this.dtype);
-		hasher.update(JSON.stringify(this.arg));
+		if (this.op === AluOp.Const) hasher.update(this.arg);
+		else hasher.update(JSON.stringify(this.arg));
 		hasher.update(this.src.length);
 		for (const s of this.src) hasher.update(s);
 		this.#hash = hasher.value;
@@ -757,6 +766,7 @@ var AluExp = class AluExp {
 			if (op === AluOp.Mul && x === 1) return src[1 - i];
 			if (op === AluOp.Mul && x === 0) return AluExp.const(this.dtype, 0);
 			if (op === AluOp.Idiv && i === 1 && x === 1) return src[1 - i];
+			if (op === AluOp.Cmpne && src[i].dtype === DType.Bool && x === 0) return src[1 - i];
 		}
 		if ((op === AluOp.Add || op === AluOp.Sub) && src[1].op === AluOp.Mul) {
 			const [a, b] = src[1].src;
@@ -958,11 +968,13 @@ var AluExp = class AluExp {
 					else if (fromType === DType.Int32) view.setInt32(0, x, true);
 					else if (fromType === DType.Uint32) view.setUint32(0, x, true);
 					else if (fromType === DType.Float16) view.setFloat16(0, x, true);
+					else if (fromType === DType.Float64) view.setFloat64(0, x, true);
 					else throw new Error(`Unsupported bitcast from ${fromType}`);
 					if (this.dtype === DType.Float32) return view.getFloat32(0, true);
 					else if (this.dtype === DType.Int32) return view.getInt32(0, true);
 					else if (this.dtype === DType.Uint32) return view.getUint32(0, true);
 					else if (this.dtype === DType.Float16) return view.getFloat16(0, true);
+					else if (this.dtype === DType.Float64) return view.getFloat64(0, true);
 					else throw new Error(`Unsupported bitcast to ${this.dtype}`);
 				}
 				default: throw new Error(`Missing implemementation for ${this.op}`);
@@ -2925,6 +2937,7 @@ var CodeGenerator = class {
 	local;
 	i32;
 	f32;
+	f64;
 	v128;
 	i32x4;
 	f32x4;
@@ -2944,6 +2957,7 @@ var CodeGenerator = class {
 		this.local = new Local(this);
 		this.i32 = new I32(this);
 		this.f32 = new F32(this);
+		this.f64 = new F64(this);
 		this.v128 = new V128(this);
 		this.i32x4 = new I32x4(this);
 		this.f32x4 = new F32x4(this);
@@ -3330,6 +3344,8 @@ var I32 = class {
 	ne = BINARY_OP("ne", 71, "i32", "i32", "i32");
 	trunc_f32_s = UNARY_OP("trunc_f32_s", 168, "f32", "i32");
 	trunc_f32_u = UNARY_OP("trunc_f32_u", 169, "f32", "i32");
+	trunc_f64_s = UNARY_OP("trunc_f64_s", 170, "f64", "i32");
+	trunc_f64_u = UNARY_OP("trunc_f64_u", 171, "f64", "i32");
 	load = LOAD_OP("load", 40, "i32");
 	load8_s = LOAD_OP("load8_s", 44, "i32");
 	load8_u = LOAD_OP("load8_u", 45, "i32");
@@ -3341,6 +3357,8 @@ var I32 = class {
 	reinterpret_f32 = UNARY_OP("reinterpret_f32", 188, "f32", "i32");
 	trunc_sat_f32_s = UNARY_OP("trunc_sat_f32_s", [252, 0], "f32", "i32");
 	trunc_sat_f32_u = UNARY_OP("trunc_sat_f32_u", [252, 1], "f32", "i32");
+	trunc_sat_f64_s = UNARY_OP("trunc_sat_f64_s", [252, 2], "f64", "i32");
+	trunc_sat_f64_u = UNARY_OP("trunc_sat_f64_u", [252, 3], "f64", "i32");
 };
 var F32 = class {
 	constructor(cg) {
@@ -3360,6 +3378,8 @@ var F32 = class {
 		for (let i = 0; i < 4; i++) this.cg._emit(bytes[i]);
 		this.cg._push(this);
 	}
+	load = LOAD_OP("load", 42, "f32");
+	store = STORE_OP("store", 56, "f32");
 	eq = BINARY_OP("eq", 91, "f32", "f32", "i32");
 	ne = BINARY_OP("ne", 92, "f32", "f32", "i32");
 	lt = BINARY_OP("lt", 93, "f32", "f32", "i32");
@@ -3382,10 +3402,53 @@ var F32 = class {
 	copysign = BINARY_OP("copysign", 152, "f32", "f32", "f32");
 	convert_i32_s = UNARY_OP("convert_i32_s", 178, "i32", "f32");
 	convert_i32_u = UNARY_OP("convert_i32_u", 179, "i32", "f32");
-	load = LOAD_OP("load", 42, "f32");
-	store = STORE_OP("store", 56, "f32");
+	demote_f64 = UNARY_OP("demote_f64", 182, "f64", "f32");
 	reinterpret_i32 = UNARY_OP("reinterpret_i32", 190, "i32", "f32");
 };
+var F64 = class {
+	constructor(cg) {
+		this.cg = cg;
+	}
+	get typeId() {
+		return 124;
+	}
+	get name() {
+		return "f64";
+	}
+	const(f) {
+		this.cg._emit(68);
+		const buffer = /* @__PURE__ */ new ArrayBuffer(8);
+		new DataView(buffer).setFloat64(0, f, true);
+		const bytes = new Uint8Array(buffer);
+		for (let i = 0; i < 8; i++) this.cg._emit(bytes[i]);
+		this.cg._push(this);
+	}
+	load = LOAD_OP("load", 43, "f64");
+	store = STORE_OP("store", 57, "f64");
+	eq = BINARY_OP("eq", 97, "f64", "f64", "i32");
+	ne = BINARY_OP("ne", 98, "f64", "f64", "i32");
+	lt = BINARY_OP("lt", 99, "f64", "f64", "i32");
+	gt = BINARY_OP("gt", 100, "f64", "f64", "i32");
+	le = BINARY_OP("le", 101, "f64", "f64", "i32");
+	ge = BINARY_OP("ge", 102, "f64", "f64", "i32");
+	abs = UNARY_OP("abs", 153, "f64", "f64");
+	neg = UNARY_OP("neg", 154, "f64", "f64");
+	ceil = UNARY_OP("ceil", 155, "f64", "f64");
+	floor = UNARY_OP("floor", 156, "f64", "f64");
+	trunc = UNARY_OP("trunc", 157, "f64", "f64");
+	nearest = UNARY_OP("nearest", 158, "f64", "f64");
+	sqrt = UNARY_OP("sqrt", 159, "f64", "f64");
+	add = BINARY_OP("add", 160, "f64", "f64", "f64");
+	sub = BINARY_OP("sub", 161, "f64", "f64", "f64");
+	mul = BINARY_OP("mul", 162, "f64", "f64", "f64");
+	div = BINARY_OP("div", 163, "f64", "f64", "f64");
+	min = BINARY_OP("min", 164, "f64", "f64", "f64");
+	max = BINARY_OP("max", 165, "f64", "f64", "f64");
+	copysign = BINARY_OP("copysign", 166, "f64", "f64", "f64");
+	convert_i32_s = UNARY_OP("convert_i32_s", 183, "i32", "f64");
+	convert_i32_u = UNARY_OP("convert_i32_u", 184, "i32", "f64");
+	promote_f32 = UNARY_OP("promote_f32", 187, "f32", "f64");
+};
 function VECTOR_OP(op, vopcode, inTypes, outType) {
 	return function() {
 		for (const inType of inTypes.toReversed()) {
@@ -3638,10 +3701,10 @@ function codegenWasm(kernel) {
 				cg.local.get(acc);
 				if (re.dtype === DType.Bool) cg.i32.and();
 				else dty(cg, re.op, re.dtype).mul();
-			} else if (re.op === AluOp.Min || re.op === AluOp.Max) if (re.dtype === DType.Float32) {
+			} else if (re.op === AluOp.Min || re.op === AluOp.Max) if (isFloatDtype(re.dtype)) {
 				cg.local.get(acc);
-				if (re.op === AluOp.Min) cg.f32.min();
-				else cg.f32.max();
+				if (re.op === AluOp.Min) dtyF(cg, re.op, re.dtype).min();
+				else dtyF(cg, re.op, re.dtype).max();
 			} else if ([
 				DType.Int32,
 				DType.Uint32,
@@ -3703,27 +3766,30 @@ function translateExp(cg, funcs, exp, ctx) {
 			else if (op === AluOp.Sub) dty(cg, op, dtype).sub();
 			else if (op === AluOp.Mul) if (dtype === DType.Bool) cg.i32.and();
 			else dty(cg, op, dtype).mul();
-			else if (op === AluOp.Idiv) if (dtype === DType.Float32) cg.f32.div(), cg.f32.trunc();
-			else if (dtype === DType.Uint32) cg.i32.div_u();
+			else if (op === AluOp.Idiv) if (isFloatDtype(dtype)) {
+				dtyF(cg, op, dtype).div();
+				dtyF(cg, op, dtype).trunc();
+			} else if (dtype === DType.Uint32) cg.i32.div_u();
 			else if (dtype === DType.Int32) cg.i32.div_s();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-			else if (op === AluOp.Mod) if (dtype === DType.Float32) {
-				const a = cg.local.declare(cg.f32);
-				const b = cg.local.declare(cg.f32);
+			else if (op === AluOp.Mod) if (isFloatDtype(dtype)) {
+				const dt = dtyF(cg, op, dtype);
+				const a = cg.local.declare(dt);
+				const b = cg.local.declare(dt);
 				cg.local.set(b);
 				cg.local.tee(a);
 				cg.local.get(a);
 				cg.local.get(b);
-				cg.f32.div();
-				cg.f32.trunc();
+				dt.div();
+				dt.trunc();
 				cg.local.get(b);
-				cg.f32.mul();
-				cg.f32.sub();
+				dt.mul();
+				dt.sub();
 			} else if (dtype === DType.Uint32) cg.i32.rem_u();
 			else if (dtype === DType.Int32) cg.i32.rem_s();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-			else if (op === AluOp.Min || op === AluOp.Max) if (dtype === DType.Float32) if (op === AluOp.Min) cg.f32.min();
-			else cg.f32.max();
+			else if (op === AluOp.Min || op === AluOp.Max) if (isFloatDtype(dtype)) if (op === AluOp.Min) dtyF(cg, op, dtype).min();
+			else dtyF(cg, op, dtype).max();
 			else if (dtype === DType.Int32 || dtype === DType.Uint32) {
 				const a = cg.local.declare(cg.i32);
 				const b = cg.local.declare(cg.i32);
@@ -3740,54 +3806,74 @@ function translateExp(cg, funcs, exp, ctx) {
 			} else throw new UnsupportedOpError(op, dtype, "wasm");
 			else if (op === AluOp.Cmplt) {
 				const srcDtype = src[0].dtype;
-				if (srcDtype === DType.Float32) cg.f32.lt();
+				if (isFloatDtype(srcDtype)) dtyF(cg, op, srcDtype).lt();
 				else if (srcDtype === DType.Int32) cg.i32.lt_s();
 				else if (srcDtype === DType.Uint32) cg.i32.lt_u();
 				else throw new UnsupportedOpError(op, dtype, "wasm");
 			} else if (op === AluOp.Cmpne) dty(cg, op, src[0].dtype).ne();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else if (AluGroup.Unary.has(op)) if (op === AluOp.Sin) gen(src[0]), cg.call(funcs.sin);
-		else if (op === AluOp.Cos) gen(src[0]), cg.call(funcs.cos);
-		else if (op === AluOp.Asin) gen(src[0]), cg.call(funcs.asin);
-		else if (op === AluOp.Atan) gen(src[0]), cg.call(funcs.atan);
-		else if (op === AluOp.Exp) gen(src[0]), cg.call(funcs.exp);
-		else if (op === AluOp.Log) gen(src[0]), cg.call(funcs.log);
-		else if (op === AluOp.Erf) gen(src[0]), cg.call(funcs.erf);
-		else if (op === AluOp.Erfc) gen(src[0]), cg.call(funcs.erfc);
-		else if (op === AluOp.Sqrt) gen(src[0]), cg.f32.sqrt();
-		else if (op === AluOp.Reciprocal) cg.f32.const(1), gen(src[0]), cg.f32.div();
-		else if (op === AluOp.Cast) {
-			gen(src[0]);
-			const dtype0 = src[0].dtype;
-			const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
-			if (dtype === DType.Int32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_s();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_u();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Float32) if (dtype0 === DType.Float32);
-			else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f32.convert_i32_s();
-			else if (dtype0 === DType.Uint32) cg.f32.convert_i32_u();
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Bool) if (dtype0 === DType.Bool);
-			else if (i32repr) cg.i32.const(0), cg.i32.ne();
-			else if (dtype0 === DType.Float32) cg.f32.const(0), cg.f32.ne();
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else if (op === AluOp.Bitcast) {
-			gen(src[0]);
-			const dtype0 = src[0].dtype;
-			const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32;
-			if (dtype === DType.Int32 || dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.reinterpret_f32();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Float32) if (i32repr) cg.f32.reinterpret_i32();
-			else if (dtype0 === DType.Float32);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else throw new UnsupportedOpError(op, dtype, "wasm");
-		else if (op === AluOp.Where) {
+		} else if (AluGroup.Unary.has(op)) {
+			const callFuncF32 = (func) => {
+				if (dtype !== DType.Float32) if (dtype === DType.Float64) cg.f32.demote_f64();
+				else throw new UnsupportedOpError(op, dtype, "wasm");
+				cg.call(func);
+				if (dtype === DType.Float64) cg.f64.promote_f32();
+			};
+			if (op === AluOp.Sin) gen(src[0]), callFuncF32(funcs.sin);
+			else if (op === AluOp.Cos) gen(src[0]), callFuncF32(funcs.cos);
+			else if (op === AluOp.Asin) gen(src[0]), callFuncF32(funcs.asin);
+			else if (op === AluOp.Atan) gen(src[0]), callFuncF32(funcs.atan);
+			else if (op === AluOp.Exp) gen(src[0]), callFuncF32(funcs.exp);
+			else if (op === AluOp.Log) gen(src[0]), callFuncF32(funcs.log);
+			else if (op === AluOp.Erf) gen(src[0]), callFuncF32(funcs.erf);
+			else if (op === AluOp.Erfc) gen(src[0]), callFuncF32(funcs.erfc);
+			else if (op === AluOp.Sqrt) gen(src[0]), dtyF(cg, op, dtype).sqrt();
+			else if (op === AluOp.Reciprocal) {
+				const dt = dtyF(cg, op, dtype);
+				dt.const(1), gen(src[0]), dt.div();
+			} else if (op === AluOp.Cast) {
+				gen(src[0]);
+				const dtype0 = src[0].dtype;
+				const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
+				if (dtype === DType.Int32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_s();
+				else if (dtype0 === DType.Float64) cg.i32.trunc_sat_f64_s();
+				else if (i32repr);
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_u();
+				else if (dtype0 === DType.Float64) cg.i32.trunc_sat_f64_u();
+				else if (i32repr);
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Float32) if (dtype0 === DType.Float32);
+				else if (dtype0 === DType.Float64) cg.f32.demote_f64();
+				else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f32.convert_i32_s();
+				else if (dtype0 === DType.Uint32) cg.f32.convert_i32_u();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Float64) if (dtype0 === DType.Float32) cg.f64.promote_f32();
+				else if (dtype0 === DType.Float64);
+				else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f64.convert_i32_s();
+				else if (dtype0 === DType.Uint32) cg.f64.convert_i32_u();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Bool) if (dtype0 === DType.Bool);
+				else if (i32repr) cg.i32.const(0), cg.i32.ne();
+				else if (dtype0 === DType.Float32) cg.f32.const(0), cg.f32.ne();
+				else if (dtype0 === DType.Float64) cg.f64.const(0), cg.f64.ne();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else throw new UnsupportedOpError(op, dtype, "wasm");
+			} else if (op === AluOp.Bitcast) {
+				gen(src[0]);
+				const dtype0 = src[0].dtype;
+				if (dtype !== dtype0) {
+					const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32;
+					if (dtype === DType.Int32 || dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.reinterpret_f32();
+					else if (i32repr);
+					else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+					else if (dtype === DType.Float32) if (i32repr) cg.f32.reinterpret_i32();
+					else if (dtype0 === DType.Float32);
+					else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+					else throw new UnsupportedOpError(op, dtype, "wasm");
+				}
+			} else throw new UnsupportedOpError(op, dtype, "wasm");
+		} else if (op === AluOp.Where) {
 			gen(src[1]);
 			gen(src[2]);
 			gen(src[0]);
@@ -3832,12 +3918,20 @@ function translateExp(cg, funcs, exp, ctx) {
 function dty(cg, op, dtype) {
 	switch (dtype) {
 		case DType.Float32: return cg.f32;
+		case DType.Float64: return cg.f64;
 		case DType.Int32:
 		case DType.Uint32:
 		case DType.Bool: return cg.i32;
 		default: throw new UnsupportedOpError(op, dtype, "wasm");
 	}
 }
+function dtyF(cg, op, dtype) {
+	switch (dtype) {
+		case DType.Float32: return cg.f32;
+		case DType.Float64: return cg.f64;
+		default: throw new UnsupportedOpError(op, dtype, "wasm");
+	}
+}
 
 //#endregion
 //#region src/backend.ts
@@ -3846,10 +3940,10 @@ const devices = [
 	"wasm",
 	"webgpu"
 ];
-let defaultBackend = "wasm";
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
-initializedBackends.set("wasm", new WasmBackend());
+if (typeof WebAssembly !== "undefined") initializedBackends.set("wasm", new WasmBackend());
+let defaultBackend = initializedBackends.has("wasm") ? "wasm" : "cpu";
 /** Configure the default device for arrays. */
 function defaultDevice(device) {
 	if (device !== void 0) if (initializedBackends.has(device)) defaultBackend = device;
@@ -3876,12 +3970,14 @@ async function init(...devicesToInit) {
 /** Create a backend, if available. Internal function called by `init()`. */
 async function createBackend(device) {
 	if (device === "cpu") return new CpuBackend();
-	else if (device === "wasm") return new WasmBackend();
-	else if (device === "webgpu") {
+	else if (device === "wasm") {
+		if (typeof WebAssembly === "undefined") return null;
+		return new WasmBackend();
+	} else if (device === "webgpu") {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-LGi2A3mS.js");
+		const { WebGPUBackend } = await import("./webgpu-B3UVme6n.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -3935,4 +4031,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-DwIAd0AG.js.map
+//# sourceMappingURL=backend-CoVtc9dx.js.map

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__ge
 }) : target, mod));
 
 //#endregion
-const require_backend = require('./backend-FtkbO6pI.cjs');
+const require_backend = require('./backend-BbrKEB18.cjs');
 
 //#region src/tree.ts
 var tree_exports = {};
@@ -401,11 +401,9 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	return Primitive$1;
 }({});
 let CompareOp = /* @__PURE__ */ function(CompareOp$1) {
-	CompareOp$1["Greater"] = "greater";
 	CompareOp$1["Less"] = "less";
 	CompareOp$1["Equal"] = "equal";
 	CompareOp$1["NotEqual"] = "not_equal";
-	CompareOp$1["GreaterEqual"] = "greater_equal";
 	CompareOp$1["LessEqual"] = "less_equal";
 	return CompareOp$1;
 }({});
@@ -501,7 +499,7 @@ function compare(x, y, op) {
 	return bind1(Primitive.Compare, [x, y], { op });
 }
 function greater$1(x, y) {
-	return compare(x, y, CompareOp.Greater);
+	return compare(y, x, CompareOp.Less);
 }
 function less$1(x, y) {
 	return compare(x, y, CompareOp.Less);
@@ -513,7 +511,7 @@ function notEqual$1(x, y) {
 	return compare(x, y, CompareOp.NotEqual);
 }
 function greaterEqual$1(x, y) {
-	return compare(x, y, CompareOp.GreaterEqual);
+	return compare(y, x, CompareOp.LessEqual);
 }
 function lessEqual$1(x, y) {
 	return compare(x, y, CompareOp.LessEqual);
@@ -2240,6 +2238,9 @@ function arrayFromData(data, shape$1, { dtype, device }, weakType = false) {
 	} else if (data instanceof Float16Array) {
 		if (dtype && dtype !== require_backend.DType.Float16) throw new Error("Float16Array must have float16 type");
 		dtype ??= require_backend.DType.Float16;
+	} else if (data instanceof Float64Array) {
+		if (dtype && dtype !== require_backend.DType.Float64) throw new Error("Float64Array must have float64 type");
+		dtype ??= require_backend.DType.Float64;
 	} else throw new Error("Unsupported data array type: " + data.constructor.name);
 	if (data.length < inlineArrayLimit) {
 		let allEqual = true;
@@ -2451,11 +2452,9 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 }
 function aluCompare(a, b, op) {
 	switch (op) {
-		case CompareOp.Greater: return require_backend.AluExp.mul(require_backend.AluExp.cmpne(a, b), require_backend.AluExp.cmplt(a, b).not());
 		case CompareOp.Less: return require_backend.AluExp.cmplt(a, b);
 		case CompareOp.Equal: return require_backend.AluExp.cmpne(a, b).not();
 		case CompareOp.NotEqual: return require_backend.AluExp.cmpne(a, b);
-		case CompareOp.GreaterEqual: return require_backend.AluExp.cmplt(a, b).not();
 		case CompareOp.LessEqual: return require_backend.AluExp.add(require_backend.AluExp.cmplt(a, b), require_backend.AluExp.cmpne(a, b).not());
 	}
 }
@@ -2592,7 +2591,7 @@ var JaxprEqn = class {
 		const paramsList = Object.entries(this.params).map(([k, v]) => require_backend.PPrint.pp(`${k}=${v}`));
 		if (paramsList.length > 0) rhs = rhs.stack(require_backend.PPrint.pp(" [ ")).stack(require_backend.PPrint.prototype.concat(...paramsList)).stack(require_backend.PPrint.pp(" ] "));
 		else rhs = rhs.stack(require_backend.PPrint.pp(" "));
-		rhs = rhs.stack(require_backend.PPrint.pp(this.inputs.map((x) => x instanceof Var ? vp.name(x) : JSON.stringify(x.value)).join(" ")));
+		rhs = rhs.stack(require_backend.PPrint.pp(this.inputs.map((x) => x instanceof Var ? vp.name(x) : String(x.value)).join(" ")));
 		return lhs.stack(require_backend.PPrint.pp(" = ")).stack(rhs);
 	}
 	toString() {
@@ -4379,6 +4378,7 @@ __export(numpy_exports, {
 	flipud: () => flipud,
 	float16: () => float16,
 	float32: () => float32,
+	float64: () => float64,
 	full: () => full,
 	fullLike: () => fullLike$1,
 	greater: () => greater,
@@ -4392,6 +4392,11 @@ __export(numpy_exports, {
 	inf: () => inf,
 	inner: () => inner,
 	int32: () => int32,
+	isfinite: () => isfinite,
+	isinf: () => isinf,
+	isnan: () => isnan,
+	isneginf: () => isneginf,
+	isposinf: () => isposinf,
 	less: () => less,
 	lessEqual: () => lessEqual,
 	linspace: () => linspace,
@@ -4462,6 +4467,7 @@ const int32 = require_backend.DType.Int32;
 const uint32 = require_backend.DType.Uint32;
 const bool = require_backend.DType.Bool;
 const float16 = require_backend.DType.Float16;
+const float64 = require_backend.DType.Float64;
 /** Euler's constant, `e = 2.7182818284590...` */
 const e = Math.E;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -5262,6 +5268,34 @@ function var_(x, axis = null, opts) {
 function std(x, axis = null, opts) {
 	return sqrt(var_(x, axis, opts));
 }
+/** Test element-wise for positive or negative infinity, return bool array. */
+function isinf(x) {
+	x = fudgeArray(x);
+	return require_backend.isFloatDtype(x.dtype) ? x.ref.equal(Infinity).add(x.equal(-Infinity)) : fullLike$1(x, false);
+}
+/** Test element-wise for NaN (Not a Number). */
+function isnan(x) {
+	x = fudgeArray(x);
+	return require_backend.isFloatDtype(x.dtype) ? x.ref.notEqual(x) : fullLike$1(x, false);
+}
+/** Test element-wise for negative infinity, return bool array. */
+function isneginf(x) {
+	x = fudgeArray(x);
+	return require_backend.isFloatDtype(x.dtype) ? x.equal(-Infinity) : fullLike$1(x, false);
+}
+/** Test element-wise for positive infinity, return bool array. */
+function isposinf(x) {
+	x = fudgeArray(x);
+	return require_backend.isFloatDtype(x.dtype) ? x.equal(Infinity) : fullLike$1(x, false);
+}
+/**
+* @function
+* Test element-wise for finite values (not infinity or NaN).
+*/
+const isfinite = jit$1(function isfinite$1(x) {
+	if (!require_backend.isFloatDtype(x.dtype)) return fullLike$1(x, true);
+	return isnan(x.ref).add(isinf(x)).notEqual(true);
+});
 
 //#endregion
 //#region src/nn.ts

package/dist/index.d.cts

@@ -168,9 +168,10 @@ declare enum DType {
   Uint32 = "uint32",
   Bool = "bool",
   Float16 = "float16",
+  Float64 = "float64",
 }
 /** @inline */
-type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer>;
+type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer> | Float64Array<ArrayBuffer>;
 /**
  * Promote two dtypes to their join according to the type lattice.
  *
@@ -180,7 +181,7 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  *
  * **Type lattice:**
  * ```text
- * bool -> uint32 -> int32 -> float16 -> float32
+ * bool -> uint32 -> int32 -> float16 -> float32 -> float64
  *  weakType --^
  * ```
  *
@@ -243,6 +244,7 @@ declare class AluExp implements FpHashable {
   static u32(value: number): AluExp;
   static bool(value: boolean): AluExp;
   static f16(value: number): AluExp;
+  static f64(value: number): AluExp;
   not(): AluExp;
   /** Compute a reasonable expression hash with low collision rate. */
   getHash(): bigint;
@@ -630,11 +632,9 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
 /** Type of parameters taken by each primitive. */
 type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
 declare enum CompareOp {
-  Greater = "greater",
   Less = "less",
   Equal = "equal",
   NotEqual = "not_equal",
-  GreaterEqual = "greater_equal",
   LessEqual = "less_equal",
 }
 /** @inline */
@@ -985,7 +985,7 @@ declare class Array extends Tracer {
   _putSync(backend: Backend): Array;
 }
 /** Constructor for creating a new array from data. */
-declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
+declare function array(values: Array | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
   dtype,
   device
@@ -1058,13 +1058,14 @@ declare function linspace(start: number, stop: number, num?: number, endpoint?:
   device
 }?: DTypeAndDevice): Array;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, clip, columnStack, concatenate, cos, cosh, deg2rad, degrees, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, pow, power, prod, promoteTypes, rad2deg, radians, ravel, reciprocal, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, stack, std, subtract, sum, tan, tanh, tile, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, clip, columnStack, concatenate, cos, cosh, deg2rad, degrees, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, pow, power, prod, promoteTypes, rad2deg, radians, ravel, reciprocal, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, stack, std, subtract, sum, tan, tanh, tile, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
 declare const uint32 = DType.Uint32;
 declare const bool = DType.Bool;
 declare const float16 = DType.Float16;
+declare const float64 = DType.Float64;
 /** Euler's constant, `e = 2.7182818284590...` */
 declare const e: number;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -1535,6 +1536,19 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+/** Test element-wise for positive or negative infinity, return bool array. */
+declare function isinf(x: ArrayLike): Array;
+/** Test element-wise for NaN (Not a Number). */
+declare function isnan(x: ArrayLike): Array;
+/** Test element-wise for negative infinity, return bool array. */
+declare function isneginf(x: ArrayLike): Array;
+/** Test element-wise for positive infinity, return bool array. */
+declare function isposinf(x: ArrayLike): Array;
+/**
+ * @function
+ * Test element-wise for finite values (not infinity or NaN).
+ */
+declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
 //# sourceMappingURL=numpy.d.ts.map
 //#endregion
 //#region src/frontend/jaxpr.d.ts