@jax-js/jax 0.1.10 → 0.1.11

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, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-Ctqs8la1.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, getWebGPUDevice, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DZvR7mZV.js";
 
 //#region src/frontend/convolution.ts
 /**
@@ -333,6 +333,8 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	Primitive$1["Mod"] = "mod";
 	Primitive$1["Min"] = "min";
 	Primitive$1["Max"] = "max";
+	Primitive$1["BitCombine"] = "bit_combine";
+	Primitive$1["BitShift"] = "bit_shift";
 	Primitive$1["Neg"] = "neg";
 	Primitive$1["Reciprocal"] = "reciprocal";
 	Primitive$1["Floor"] = "floor";
@@ -406,6 +408,12 @@ function min$1(x, y) {
 function max$1(x, y) {
 	return bind1(Primitive.Max, [x, y]);
 }
+function bitCombine(x, y, op) {
+	return bind1(Primitive.BitCombine, [x, y], { op });
+}
+function bitShift(x, y, op) {
+	return bind1(Primitive.BitShift, [x, y], { op });
+}
 function neg(x) {
 	return bind1(Primitive.Neg, [x]);
 }
@@ -1620,6 +1628,16 @@ const abstractEvalRules = {
 	[Primitive.Mod]: binopAbstractEval,
 	[Primitive.Min]: binopAbstractEval,
 	[Primitive.Max]: binopAbstractEval,
+	[Primitive.BitCombine]([x, y]) {
+		const aval = promoteAvals(x, y);
+		if (isFloatDtype(aval.dtype)) throw new TypeError(`bitwise operations require integer or boolean inputs, got ${aval.dtype}`);
+		return [aval];
+	},
+	[Primitive.BitShift]([x, y]) {
+		const shape$1 = generalBroadcast(x.shape, y.shape);
+		if (isFloatDtype(x.dtype) || isFloatDtype(y.dtype) || x.dtype === DType.Bool || y.dtype === DType.Bool) throw new TypeError(`bit shift operations require integer inputs, got ${x} and ${y}`);
+		return [new ShapedArray(shape$1, x.dtype, x.weakType)];
+	},
 	[Primitive.Neg]: vectorizedUnopAbstractEval,
 	[Primitive.Reciprocal]: vectorizedUnopAbstractEval,
 	[Primitive.Floor]: vectorizedUnopAbstractEval,
@@ -2155,6 +2173,8 @@ const jitRules = {
 	[Primitive.Mod]: broadcastedJit(([a, b]) => AluExp.mod(a, b)),
 	[Primitive.Min]: broadcastedJit(([a, b]) => AluExp.min(a, b)),
 	[Primitive.Max]: broadcastedJit(([a, b]) => AluExp.max(a, b)),
+	[Primitive.BitCombine]: broadcastedJit(([a, b], { op }) => AluExp.bitCombine(a, b, op)),
+	[Primitive.BitShift]: broadcastedJit(([a, b], { op }) => AluExp.bitShift(a, b, op)),
 	[Primitive.Neg]: unopJit((a) => AluExp.sub(AluExp.const(a.dtype, 0), a)),
 	[Primitive.Reciprocal]: unopJit(AluExp.reciprocal),
 	[Primitive.Floor]: unopJit(AluExp.floor),
@@ -2347,7 +2367,9 @@ function splitGraphDataflow(backend, jaxpr) {
 					case Primitive.Idiv:
 					case Primitive.Mod:
 					case Primitive.Min:
-					case Primitive.Max: {
+					case Primitive.Max:
+					case Primitive.BitCombine:
+					case Primitive.BitShift: {
 						const otherInput = nextEqn.inputs.find((v) => v !== outVar);
 						if (otherInput instanceof Lit || deepEqual(generalBroadcast(otherInput.aval.shape, outVar.aval.shape), outVar.aval.shape)) {
 							head = usages[0];
@@ -2986,6 +3008,42 @@ var Array$1 = class Array$1 extends Tracer {
 		return dtypedArray(this.dtype, buf);
 	}
 	/**
+	* Return this array as a WebGPU buffer (with `STORAGE | COPY_SRC`).
+	*
+	* Only available on the WebGPU backend. The array's memory is still managed
+	* by jax-js, and it will be freed when the buffer is no longer in use. You
+	* _should not_ mutate the buffer's contents.
+	*
+	* Note that the GPU buffer may be slightly larger than the array's size; it
+	* will always be aligned to 4 bytes.
+	*/
+	async gpuBuffer() {
+		if (this.device !== "webgpu") throw new Error(`gpuBuffer() is only available on WebGPU backend`);
+		this.#realize();
+		const pending = this.#pending;
+		if (pending) {
+			await Promise.all(pending.map((p) => p.prepare()));
+			for (const p of pending) p.submit();
+		}
+		const backend = this.#backend;
+		const { buffer } = backend.buffers.get(this.#source);
+		this.dispose();
+		return buffer;
+	}
+	/** Synchronous version of `Array.gpuBuffer()`. */
+	gpuBufferSync() {
+		if (this.device !== "webgpu") throw new Error(`gpuBufferSync() is only available on WebGPU backend`);
+		this.#realize();
+		for (const p of this.#pending) {
+			p.prepareSync();
+			p.submit();
+		}
+		const backend = this.#backend;
+		const { buffer } = backend.buffers.get(this.#source);
+		this.dispose();
+		return buffer;
+	}
+	/**
 	* Convert this array into a JavaScript object.
 	*
 	* This is a blocking operation that will compile all of the shaders and wait
@@ -3032,6 +3090,14 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Max]([x, y]) {
 				return [x.#binary(AluOp.Max, y)];
 			},
+			[Primitive.BitCombine]([x, y], { op }) {
+				const custom = (src) => AluExp.bitCombine(src[0], src[1], op);
+				return [Array$1.#naryCustom("bit_combine", custom, [x, y])];
+			},
+			[Primitive.BitShift]([x, y], { op }) {
+				const custom = (src) => AluExp.bitShift(src[0], src[1], op);
+				return [Array$1.#naryCustom("bit_shift", custom, [x, y], { dtypeOverride: [void 0, y.dtype] })];
+			},
 			[Primitive.Neg]([x]) {
 				return [zerosLike$1(x.ref).#binary(AluOp.Sub, x)];
 			},
@@ -3723,6 +3789,8 @@ const vmapRules = {
 	[Primitive.Mod]: broadcastBatcher(Primitive.Mod),
 	[Primitive.Min]: broadcastBatcher(Primitive.Min),
 	[Primitive.Max]: broadcastBatcher(Primitive.Max),
+	[Primitive.BitCombine]: broadcastBatcher(Primitive.BitCombine),
+	[Primitive.BitShift]: broadcastBatcher(Primitive.BitShift),
 	[Primitive.Neg]: unopBatcher(Primitive.Neg),
 	[Primitive.Reciprocal]: unopBatcher(Primitive.Reciprocal),
 	[Primitive.Floor]: unopBatcher(Primitive.Floor),
@@ -4045,6 +4113,8 @@ const jvpRules = {
 	[Primitive.Max]([x, y], [dx, dy]) {
 		return [[max$1(x.ref, y.ref)], [where$1(less$1(x, y), dy, dx)]];
 	},
+	[Primitive.BitCombine]: zeroTangentsJvp(Primitive.BitCombine),
+	[Primitive.BitShift]: zeroTangentsJvp(Primitive.BitShift),
 	[Primitive.Neg]: linearTangentsJvp(Primitive.Neg),
 	[Primitive.Reciprocal]([x], [dx]) {
 		const xRecip = reciprocal$1(x.ref);
@@ -5236,7 +5306,8 @@ __export(numpy_linalg_exports, {
 	solve: () => solve,
 	tensordot: () => tensordot,
 	trace: () => trace,
-	vecdot: () => vecdot
+	vecdot: () => vecdot,
+	vectorNorm: () => vectorNorm
 });
 function checkSquare(name, a) {
 	if (a.ndim < 2 || a.shape[a.ndim - 1] !== a.shape[a.ndim - 2]) throw new Error(`${name}: input must be at least 2D square matrix, got ${a.aval}`);
@@ -5415,6 +5486,23 @@ function solve(a, b) {
 	if (bIs1d) x = squeeze(x, -1);
 	return x;
 }
+/**
+* Compute the vector norm of an array.
+*
+* @param x - Input array.
+* @param ord - Order of the norm (default 2). Supports `Infinity`, `-Infinity`, `0`, or any real number.
+* @param axis - Axis/axes to reduce over (default: all axes).
+* @param keepdims - Whether to keep reduced dimensions as size 1.
+* @returns The norm of `x`, reduced over the given axes.
+*/
+function vectorNorm(x, { ord = 2, axis = null, keepdims = false } = {}) {
+	x = fudgeArray(x);
+	const ax = axis ?? null;
+	if (ord === Infinity) return max(absolute(x), ax, { keepdims });
+	else if (ord === -Infinity) return min(absolute(x), ax, { keepdims });
+	else if (ord === 0) return x.notEqual(0).astype(x.dtype).sum(ax, { keepdims });
+	else return power(power(absolute(x), ord).sum(ax, { keepdims }), 1 / ord);
+}
 
 //#endregion
 //#region src/library/numpy/dtype-info.ts
@@ -5534,6 +5622,13 @@ __export(numpy_exports, {
 	atan2: () => atan2,
 	atanh: () => arctanh,
 	average: () => average,
+	bitwiseAnd: () => bitwiseAnd,
+	bitwiseInvert: () => invert,
+	bitwiseLeftShift: () => leftShift,
+	bitwiseNot: () => invert,
+	bitwiseOr: () => bitwiseOr,
+	bitwiseRightShift: () => rightShift,
+	bitwiseXor: () => bitwiseXor,
 	bool: () => bool,
 	broadcastArrays: () => broadcastArrays,
 	broadcastShapes: () => broadcastShapes,
@@ -5595,12 +5690,14 @@ __export(numpy_exports, {
 	inf: () => inf,
 	inner: () => inner,
 	int32: () => int32,
+	invert: () => invert,
 	isfinite: () => isfinite,
 	isinf: () => isinf,
 	isnan: () => isnan,
 	isneginf: () => isneginf,
 	isposinf: () => isposinf,
 	ldexp: () => ldexp,
+	leftShift: () => leftShift,
 	less: () => less,
 	lessEqual: () => lessEqual,
 	linalg: () => numpy_linalg_exports,
@@ -5649,6 +5746,7 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rightShift: () => rightShift,
 	rint: () => rint,
 	round: () => round,
 	shape: () => shape,
@@ -5763,6 +5861,44 @@ function logicalXor(x, y) {
 function logicalNot(x) {
 	return notEqual(astype(x, DType.Bool), true);
 }
+/** Compute element-wise bitwise AND. */
+function bitwiseAnd(x, y) {
+	return bitCombine(x, y, "and");
+}
+/** Compute element-wise bitwise OR. */
+function bitwiseOr(x, y) {
+	return bitCombine(x, y, "or");
+}
+/** Compute element-wise bitwise XOR. */
+function bitwiseXor(x, y) {
+	return bitCombine(x, y, "xor");
+}
+/** Compute element-wise bitwise NOT (inversion). */
+function invert(x) {
+	const arr = fudgeArray(x);
+	let allOnes;
+	switch (arr.dtype) {
+		case DType.Bool:
+			allOnes = true;
+			break;
+		case DType.Uint32:
+			allOnes = 4294967295;
+			break;
+		case DType.Int32:
+			allOnes = -1;
+			break;
+		default: throw new TypeError(`invert: unsupported dtype ${arr.dtype}`);
+	}
+	return bitCombine(arr, allOnes, "xor");
+}
+/** Compute element-wise left bit shift. */
+function leftShift(x, y) {
+	return bitShift(x, y, "shl");
+}
+/** Compute element-wise right bit shift. */
+function rightShift(x, y) {
+	return bitShift(x, y, "shr");
+}
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 const where = where$1;
 /**
@@ -8291,4 +8427,4 @@ async function devicePut(x, device) {
 }
 
 //#endregion
-export { Array$1 as Array, ClosedJaxpr, DType, Jaxpr, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_exports as lax, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, profiler, random_exports as random, scipy_special_exports as scipySpecial, setDebug, tree_exports as tree, valueAndGrad, vjp, vmap };
+export { Array$1 as Array, ClosedJaxpr, DType, Jaxpr, blockUntilReady, defaultDevice, devicePut, devices, getWebGPUDevice, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_exports as lax, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, profiler, random_exports as random, scipy_special_exports as scipySpecial, setDebug, tree_exports as tree, valueAndGrad, vjp, vmap };

package/dist/{webgl-CvQ1QBX1.js → webgl-D8-14NzA.js}

@@ -1,4 +1,4 @@
-import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-Ctqs8la1.js";
+import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-DZvR7mZV.js";
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `
@@ -458,6 +458,12 @@ function generateExpression(exp, args, inputDtypes) {
 			else source = `min(${a}, ${b})`;
 			else if (op === AluOp.Max) if (dtype === DType.Bool) source = `(${a} || ${b})`;
 			else source = `max(${a}, ${b})`;
+			else if (op === AluOp.BitCombine) {
+				let infix = arg === "and" ? "&" : arg === "or" ? "|" : "^";
+				if (dtype === DType.Bool) infix = infix + infix;
+				source = `(${a} ${infix} ${b})`;
+			} else if (op === AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${b})`;
 		} else if (AluGroup.Compare.has(op)) {
 			const a = gen(src[0]);
 			const b = gen(src[1]);

package/dist/{webgl-kvVt7-T7.cjs → webgl-Ovaaa-Qx.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-DMauYnfl.cjs');
+const require_backend = require('./backend-DlYlOYqN.cjs');
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `
@@ -458,6 +458,12 @@ function generateExpression(exp, args, inputDtypes) {
 			else source = `min(${a}, ${b})`;
 			else if (op === require_backend.AluOp.Max) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
 			else source = `max(${a}, ${b})`;
+			else if (op === require_backend.AluOp.BitCombine) {
+				let infix = arg === "and" ? "&" : arg === "or" ? "|" : "^";
+				if (dtype === require_backend.DType.Bool) infix = infix + infix;
+				source = `(${a} ${infix} ${b})`;
+			} else if (op === require_backend.AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${b})`;
 		} else if (require_backend.AluGroup.Compare.has(op)) {
 			const a = gen(src[0]);
 			const b = gen(src[1]);

package/dist/{webgpu-v_W_-oKw.js → webgpu-Dg8FpYrH.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, emitTrace, findPow2, isFloatDtype, isTracing, mapSetUnion, onFlushTrace, prod, range, strip1, traceSourceInfo, tuneWebgpu } from "./backend-Ctqs8la1.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, emitTrace, findPow2, isFloatDtype, isTracing, mapSetUnion, onFlushTrace, prod, range, strip1, traceSourceInfo, tuneWebgpu } from "./backend-DZvR7mZV.js";
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -1100,6 +1100,11 @@ function pipelineSource(device, kernel) {
 			else source = `min(${strip1(a)}, ${strip1(b)})`;
 			else if (op === AluOp.Max) if (dtype === DType.Bool) source = `(${a} || ${b})`;
 			else source = `max(${strip1(a)}, ${strip1(b)})`;
+			else if (op === AluOp.BitCombine) if (arg === "and") source = `(${a} & ${b})`;
+			else if (arg === "or") source = `(${a} | ${b})`;
+			else source = dtype === DType.Bool ? `(${a} != ${b})` : `(${a} ^ ${b})`;
+			else if (op === AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${b})`;
 			else if (op === AluOp.Cmplt) source = `(${a} < ${b})`;
 			else if (op === AluOp.Cmpne) if (isFloatDtype(src[0].dtype)) {
 				const x = isGensym(a) ? a : gensym();

package/dist/{webgpu-DMSx7a6M.cjs → webgpu-uU9nnttc.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-DMauYnfl.cjs');
+const require_backend = require('./backend-DlYlOYqN.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `
@@ -1100,6 +1100,11 @@ function pipelineSource(device, kernel) {
 			else source = `min(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
 			else if (op === require_backend.AluOp.Max) if (dtype === require_backend.DType.Bool) source = `(${a} || ${b})`;
 			else source = `max(${require_backend.strip1(a)}, ${require_backend.strip1(b)})`;
+			else if (op === require_backend.AluOp.BitCombine) if (arg === "and") source = `(${a} & ${b})`;
+			else if (arg === "or") source = `(${a} | ${b})`;
+			else source = dtype === require_backend.DType.Bool ? `(${a} != ${b})` : `(${a} ^ ${b})`;
+			else if (op === require_backend.AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${b})`;
 			else if (op === require_backend.AluOp.Cmplt) source = `(${a} < ${b})`;
 			else if (op === require_backend.AluOp.Cmpne) if (require_backend.isFloatDtype(src[0].dtype)) {
 				const x = isGensym(a) ? a : gensym();

package/package.json

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