@jax-js/jax 0.1.0 → 0.1.1

package/dist/index.d.ts

@@ -165,9 +165,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.
  *
@@ -177,7 +178,7 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  *
  * **Type lattice:**
  * ```text
- * bool -> uint32 -> int32 -> float16 -> float32
+ * bool -> uint32 -> int32 -> float16 -> float32 -> float64
  *  weakType --^
  * ```
  *
@@ -240,6 +241,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;
@@ -627,11 +629,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 */
@@ -982,7 +982,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
@@ -1055,13 +1055,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...` */
@@ -1532,6 +1533,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

package/dist/index.js

@@ -1,5 +1,5 @@
 import { __export } from "./chunk-Cl8Af3a2.js";
-import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DwIAd0AG.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-CoVtc9dx.js";
 
 //#region src/tree.ts
 var tree_exports = {};
@@ -370,11 +370,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;
 }({});
@@ -470,7 +468,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);
@@ -482,7 +480,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);
@@ -2209,6 +2207,9 @@ function arrayFromData(data, shape$1, { dtype, device }, weakType = false) {
 	} else if (data instanceof Float16Array) {
 		if (dtype && dtype !== DType.Float16) throw new Error("Float16Array must have float16 type");
 		dtype ??= DType.Float16;
+	} else if (data instanceof Float64Array) {
+		if (dtype && dtype !== DType.Float64) throw new Error("Float64Array must have float64 type");
+		dtype ??= DType.Float64;
 	} else throw new Error("Unsupported data array type: " + data.constructor.name);
 	if (data.length < inlineArrayLimit) {
 		let allEqual = true;
@@ -2420,11 +2421,9 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 }
 function aluCompare(a, b, op) {
 	switch (op) {
-		case CompareOp.Greater: return AluExp.mul(AluExp.cmpne(a, b), AluExp.cmplt(a, b).not());
 		case CompareOp.Less: return AluExp.cmplt(a, b);
 		case CompareOp.Equal: return AluExp.cmpne(a, b).not();
 		case CompareOp.NotEqual: return AluExp.cmpne(a, b);
-		case CompareOp.GreaterEqual: return AluExp.cmplt(a, b).not();
 		case CompareOp.LessEqual: return AluExp.add(AluExp.cmplt(a, b), AluExp.cmpne(a, b).not());
 	}
 }
@@ -2557,7 +2556,7 @@ var JaxprEqn = class {
 		const paramsList = Object.entries(this.params).map(([k, v]) => PPrint.pp(`${k}=${v}`));
 		if (paramsList.length > 0) rhs = rhs.stack(PPrint.pp(" [ ")).stack(PPrint.prototype.concat(...paramsList)).stack(PPrint.pp(" ] "));
 		else rhs = rhs.stack(PPrint.pp(" "));
-		rhs = rhs.stack(PPrint.pp(this.inputs.map((x) => x instanceof Var ? vp.name(x) : JSON.stringify(x.value)).join(" ")));
+		rhs = rhs.stack(PPrint.pp(this.inputs.map((x) => x instanceof Var ? vp.name(x) : String(x.value)).join(" ")));
 		return lhs.stack(PPrint.pp(" = ")).stack(rhs);
 	}
 	toString() {
@@ -4342,6 +4341,7 @@ __export(numpy_exports, {
 	flipud: () => flipud,
 	float16: () => float16,
 	float32: () => float32,
+	float64: () => float64,
 	full: () => full,
 	fullLike: () => fullLike$1,
 	greater: () => greater,
@@ -4355,6 +4355,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,
@@ -4425,6 +4430,7 @@ const int32 = DType.Int32;
 const uint32 = DType.Uint32;
 const bool = DType.Bool;
 const float16 = DType.Float16;
+const float64 = DType.Float64;
 /** Euler's constant, `e = 2.7182818284590...` */
 const e = Math.E;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -5225,6 +5231,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 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 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 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 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 (!isFloatDtype(x.dtype)) return fullLike$1(x, true);
+	return isnan(x.ref).add(isinf(x)).notEqual(true);
+});
 
 //#endregion
 //#region src/nn.ts

package/dist/{webgpu-LGi2A3mS.js → webgpu-B3UVme6n.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-CoVtc9dx.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,7 +426,11 @@ 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})`;
@@ -645,4 +650,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 export { WebGPUBackend };
-//# sourceMappingURL=webgpu-LGi2A3mS.js.map
+//# sourceMappingURL=webgpu-B3UVme6n.js.map

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

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-FtkbO6pI.cjs');
+const require_backend = require('./backend-BbrKEB18.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,7 +426,11 @@ 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})`;
@@ -645,4 +650,4 @@ async function compileError(shaderModule, scope, code) {
 
 //#endregion
 exports.WebGPUBackend = WebGPUBackend;
-//# sourceMappingURL=webgpu-BE7zA_01.cjs.map
+//# sourceMappingURL=webgpu-DGYNVHma.cjs.map

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jax-js/jax",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "Numerical computing and ML in the browser",
   "keywords": [
     "machine learning",
@@ -35,6 +35,19 @@
   "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",
@@ -55,9 +68,15 @@
     "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"
@@ -73,16 +92,5 @@
       }
     ],
     "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"
   }
-}
+}