@jax-js/jax 0.1.10 → 0.1.11

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod$1, isNodeMode, target) => (target = mod$1 != null ? __create(
 }) : target, mod$1));
 
 //#endregion
-const require_backend = require('./backend-DMauYnfl.cjs');
+const require_backend = require('./backend-DlYlOYqN.cjs');
 
 //#region src/frontend/convolution.ts
 /**
@@ -364,6 +364,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";
@@ -437,6 +439,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]);
 }
@@ -1655,6 +1663,16 @@ const abstractEvalRules = {
 	[Primitive.Mod]: binopAbstractEval,
 	[Primitive.Min]: binopAbstractEval,
 	[Primitive.Max]: binopAbstractEval,
+	[Primitive.BitCombine]([x, y]) {
+		const aval = promoteAvals(x, y);
+		if (require_backend.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 = require_backend.generalBroadcast(x.shape, y.shape);
+		if (require_backend.isFloatDtype(x.dtype) || require_backend.isFloatDtype(y.dtype) || x.dtype === require_backend.DType.Bool || y.dtype === require_backend.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,
@@ -2190,6 +2208,8 @@ const jitRules = {
 	[Primitive.Mod]: broadcastedJit(([a, b]) => require_backend.AluExp.mod(a, b)),
 	[Primitive.Min]: broadcastedJit(([a, b]) => require_backend.AluExp.min(a, b)),
 	[Primitive.Max]: broadcastedJit(([a, b]) => require_backend.AluExp.max(a, b)),
+	[Primitive.BitCombine]: broadcastedJit(([a, b], { op }) => require_backend.AluExp.bitCombine(a, b, op)),
+	[Primitive.BitShift]: broadcastedJit(([a, b], { op }) => require_backend.AluExp.bitShift(a, b, op)),
 	[Primitive.Neg]: unopJit((a) => require_backend.AluExp.sub(require_backend.AluExp.const(a.dtype, 0), a)),
 	[Primitive.Reciprocal]: unopJit(require_backend.AluExp.reciprocal),
 	[Primitive.Floor]: unopJit(require_backend.AluExp.floor),
@@ -2382,7 +2402,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 || require_backend.deepEqual(require_backend.generalBroadcast(otherInput.aval.shape, outVar.aval.shape), outVar.aval.shape)) {
 							head = usages[0];
@@ -3021,6 +3043,42 @@ var Array$1 = class Array$1 extends Tracer {
 		return require_backend.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
@@ -3067,6 +3125,14 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Max]([x, y]) {
 				return [x.#binary(require_backend.AluOp.Max, y)];
 			},
+			[Primitive.BitCombine]([x, y], { op }) {
+				const custom = (src) => require_backend.AluExp.bitCombine(src[0], src[1], op);
+				return [Array$1.#naryCustom("bit_combine", custom, [x, y])];
+			},
+			[Primitive.BitShift]([x, y], { op }) {
+				const custom = (src) => require_backend.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(require_backend.AluOp.Sub, x)];
 			},
@@ -3759,6 +3825,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),
@@ -4082,6 +4150,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);
@@ -5273,7 +5343,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}`);
@@ -5452,6 +5523,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
@@ -5571,6 +5659,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,
@@ -5632,12 +5727,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,
@@ -5686,6 +5783,7 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rightShift: () => rightShift,
 	rint: () => rint,
 	round: () => round,
 	shape: () => shape,
@@ -5800,6 +5898,44 @@ function logicalXor(x, y) {
 function logicalNot(x) {
 	return notEqual(astype(x, require_backend.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 require_backend.DType.Bool:
+			allOnes = true;
+			break;
+		case require_backend.DType.Uint32:
+			allOnes = 4294967295;
+			break;
+		case require_backend.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;
 /**
@@ -8336,6 +8472,7 @@ exports.blockUntilReady = blockUntilReady;
 exports.defaultDevice = require_backend.defaultDevice;
 exports.devicePut = devicePut;
 exports.devices = require_backend.devices;
+exports.getWebGPUDevice = require_backend.getWebGPUDevice;
 exports.grad = grad;
 exports.hessian = hessian;
 exports.init = require_backend.init;

package/dist/index.d.cts

@@ -232,6 +232,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -323,6 +325,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -546,6 +553,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -719,6 +731,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -767,6 +781,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1194,6 +1214,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * 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.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1571,7 +1604,7 @@ declare function fft(a: ComplexPair, axis?: number): ComplexPair;
  */
 declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
 declare namespace numpy_linalg_d_exports {
-  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
+  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1626,6 +1659,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * 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.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1679,7 +1730,7 @@ type IInfo = Readonly<{
 /** Machine limits for integer types. */
 declare function iinfo(dtype: DType): IInfo;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rightShift, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -1747,6 +1798,18 @@ declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
 declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
 /** Compute element-wise logical NOT. */
 declare function logicalNot(x: ArrayLike): Array;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -2958,4 +3021,4 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
 //#endregion
-export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, getWebGPUDevice, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };

package/dist/index.d.ts

@@ -229,6 +229,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -320,6 +322,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -543,6 +550,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -716,6 +728,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -764,6 +778,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1191,6 +1211,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * 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.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1568,7 +1601,7 @@ declare function fft(a: ComplexPair, axis?: number): ComplexPair;
  */
 declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
 declare namespace numpy_linalg_d_exports {
-  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
+  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1623,6 +1656,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * 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.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1676,7 +1727,7 @@ type IInfo = Readonly<{
 /** Machine limits for integer types. */
 declare function iinfo(dtype: DType): IInfo;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rightShift, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -1744,6 +1795,18 @@ declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
 declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
 /** Compute element-wise logical NOT. */
 declare function logicalNot(x: ArrayLike): Array;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -2955,4 +3018,4 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
 //#endregion
-export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, getWebGPUDevice, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };