@jax-js/jax 0.1.10 → 0.1.12

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-DI-V78Rk.js";
 
 //#region src/frontend/convolution.ts
 /**
@@ -209,7 +209,7 @@ __export(tree_exports, {
 	structure: () => structure,
 	unflatten: () => unflatten
 });
-const JsArray$2 = globalThis.Array;
+const JsArray$3 = globalThis.Array;
 let NodeType = /* @__PURE__ */ function(NodeType$1) {
 	NodeType$1["Array"] = "Array";
 	NodeType$1["Object"] = "Object";
@@ -257,7 +257,7 @@ function flatten(tree) {
 	return [leaves$1, treedef];
 }
 function _flatten(tree, leaves$1) {
-	if (JsArray$2.isArray(tree)) {
+	if (JsArray$3.isArray(tree)) {
 		const childTrees = tree.map((c) => _flatten(c, leaves$1));
 		return new JsTreeDef(NodeType.Array, null, childTrees);
 	} else if (typeof tree === "object" && tree !== null && tree.constructor === Object) {
@@ -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];
@@ -2438,7 +2460,7 @@ function splitGraphDataflow(backend, jaxpr) {
 
 //#endregion
 //#region src/frontend/array.ts
-const JsArray$1 = globalThis.Array;
+const JsArray$2 = globalThis.Array;
 const inlineArrayLimit = 128;
 /** Version of pureArray with fudged types. */
 const fudgeArray = pureArray;
@@ -2878,6 +2900,15 @@ var Array$1 = class Array$1 extends Tracer {
 		this.#check();
 		const indices = unravelAlu(this.#st.shape, AluVar.gidx);
 		if (this.#source instanceof AluExp) {
+			let resolvedSource;
+			if (this.#st.contiguous && this.#st.size < inlineArrayLimit && (resolvedSource = this.#source.resolve()) !== void 0) {
+				const byteLength = this.#st.size * byteWidth(this.#dtype);
+				const initialData = new Uint8Array(byteLength);
+				dtypedArray(this.#dtype, initialData).fill(resolvedSource);
+				this.#source = this.#backend.malloc(byteLength, initialData);
+				this.#st = ShapeTracker.fromShape(this.shape);
+				return;
+			}
 			const exp$2 = accessorAluExp(this.#source, this.#st, indices);
 			const kernel = new Kernel(0, this.#st.size, exp$2);
 			const output = this.#backend.malloc(kernel.bytes);
@@ -2986,6 +3017,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 +3099,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)];
 			},
@@ -3284,7 +3359,7 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 		if (!shape$1) {
 			shape$1 = [];
 			let cur = values;
-			while (JsArray$1.isArray(cur)) {
+			while (JsArray$2.isArray(cur)) {
 				shape$1.push(cur.length);
 				cur = cur[0];
 			}
@@ -3723,6 +3798,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 +4122,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);
@@ -4162,7 +4241,7 @@ const jvpRules = {
 		return [[L], [dL]];
 	},
 	[Primitive.LU]([a], [da]) {
-		const [luMatrix, pivots, permutation] = lu$1(a);
+		const [luMatrix, pivots, permutation$1] = lu$1(a);
 		const [m, n] = a.shape.slice(-2);
 		const k = Math.min(m, n);
 		const luSliceL = sliceAxis(luMatrix.ref, -1, [0, k]);
@@ -4174,7 +4253,7 @@ const jvpRules = {
 		const uPadded = n > k ? padAxis(uUpper, -2, [0, n - k]) : uUpper;
 		const uEye = n > k ? padAxis(padAxis(eye(n - k), -1, [k, 0]), -2, [k, 0]) : zerosLike$1(uPadded.ref);
 		const U = uPadded.add(uEye);
-		const P = permutation.ref.reshape([...permutation.shape, 1]).equal(arange(m)).astype(da.dtype);
+		const P = permutation$1.ref.reshape([...permutation$1.shape, 1]).equal(arange(m)).astype(da.dtype);
 		const pda = batchMatmulT(P, mT(da));
 		const la = mT(triangularSolve$1(L.ref, mT(pda), {
 			lower: true,
@@ -4186,11 +4265,11 @@ const jvpRules = {
 		return [[
 			luMatrix,
 			pivots,
-			permutation
+			permutation$1
 		], [
 			lDot.add(uDot),
 			zerosLike$1(pivots.ref),
-			zerosLike$1(permutation.ref)
+			zerosLike$1(permutation$1.ref)
 		]];
 	},
 	[Primitive.Jit](primals, tangents, { name, jaxpr }) {
@@ -5236,7 +5315,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}`);
@@ -5271,8 +5351,8 @@ function cross$1(x1, x2, axis = -1) {
 function det(a) {
 	a = fudgeArray(a);
 	const n = checkSquare("det", a);
-	const [lu$2, pivots, permutation] = lu(a);
-	permutation.dispose();
+	const [lu$2, pivots, permutation$1] = lu(a);
+	permutation$1.dispose();
 	const parity = pivots.notEqual(arange(n)).astype(int32).sum(-1).mod(2);
 	const sign$1 = parity.mul(-2).add(1);
 	const diag$1 = lu$2.diagonal(0, -1, -2);
@@ -5361,8 +5441,8 @@ function matrixPower(a, n) {
 function slogdet(a) {
 	a = fudgeArray(a);
 	const n = checkSquare("slogdet", a);
-	const [lu$2, pivots, permutation] = lu(a);
-	permutation.dispose();
+	const [lu$2, pivots, permutation$1] = lu(a);
+	permutation$1.dispose();
 	let parity = pivots.notEqual(arange(n)).astype(int32).sum(-1);
 	const diag$1 = lu$2.diagonal(0, -1, -2);
 	parity = parity.add(diag$1.ref.less(0).astype(int32).sum(-1)).mod(2);
@@ -5400,9 +5480,9 @@ function solve(a, b) {
 		n,
 		m
 	]);
-	const [lu$2, pivots, permutation] = lu(a);
+	const [lu$2, pivots, permutation$1] = lu(a);
 	pivots.dispose();
-	const P = arange(n).equal(permutation.reshape([...permutation.shape, 1])).astype(b.dtype);
+	const P = arange(n).equal(permutation$1.reshape([...permutation$1.shape, 1])).astype(b.dtype);
 	const LPb = triangularSolve(lu$2.ref, matmul(P, b), {
 		leftSide: true,
 		lower: true,
@@ -5415,6 +5495,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 +5631,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 +5699,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 +5755,7 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rightShift: () => rightShift,
 	rint: () => rint,
 	round: () => round,
 	shape: () => shape,
@@ -5763,6 +5870,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;
 /**
@@ -7193,7 +7338,7 @@ __export(lax_exports, {
 	stopGradient: () => stopGradient$1,
 	topK: () => topK
 });
-const JsArray = globalThis.Array;
+const JsArray$1 = globalThis.Array;
 /** Elementwise bitcast an array into a new dtype. */
 function bitcastConvertType(x, newDtype) {
 	return fudgeArray(x).view(newDtype);
@@ -7380,7 +7525,7 @@ function convTransposePadding(k, s, padding) {
 	} else if (padding === "VALID") {
 		padLen = k + s - 2 + Math.max(k - s, 0);
 		pad1 = k - 1;
-	} else if (JsArray.isArray(padding)) {
+	} else if (JsArray$1.isArray(padding)) {
 		const pads = [k - 1 - padding[0], k - 1 - padding[1]];
 		pad1 = pads[0];
 		padLen = pads[0] + pads[1];
@@ -7899,19 +8044,34 @@ function dotProductAttention(query, key$1, value, opts = {}) {
 //#region src/library/random.ts
 var random_exports = {};
 __export(random_exports, {
+	ball: () => ball,
 	bernoulli: () => bernoulli,
 	bits: () => bits,
 	categorical: () => categorical,
 	cauchy: () => cauchy,
+	choice: () => choice,
+	doubleSidedMaxwell: () => doubleSidedMaxwell,
 	exponential: () => exponential,
+	geometric: () => geometric,
 	gumbel: () => gumbel,
 	key: () => key,
 	laplace: () => laplace,
+	logistic: () => logistic,
+	lognormal: () => lognormal,
+	maxwell: () => maxwell,
 	multivariateNormal: () => multivariateNormal,
 	normal: () => normal,
+	pareto: () => pareto,
+	permutation: () => permutation,
+	rademacher: () => rademacher,
+	randint: () => randint,
+	rayleigh: () => rayleigh,
 	split: () => split,
-	uniform: () => uniform
+	triangular: () => triangular,
+	uniform: () => uniform,
+	weibullMin: () => weibullMin
 });
+const JsArray = globalThis.Array;
 function validateKeyShape(key$1, scalar = false) {
 	if (key$1.ndim === 0) throw new Error("Key must have at least one dimension.");
 	if (key$1.shape[key$1.shape.length - 1] !== 2) throw new Error(`Invalid key shape: ${key$1.shape}. Expected last dimension to be 2.`);
@@ -7964,6 +8124,21 @@ const uniform = jit$1(function uniform$1(key$1, shape$1 = [], { minval = 0, maxv
 	else return rand.mul(maxval - minval).add(minval);
 }, { staticArgnums: [1, 2] });
 /**
+* @function
+* Sample points uniformly from the Euclidean unit ball in `d` dimensions.
+*
+* Only the Euclidean `p=2` case is currently supported.
+*/
+const ball = jit$1(function ball$1(key$1, d, { p = 2, shape: shape$1 = [] } = {}) {
+	if (!Number.isInteger(d) || d <= 0) throw new Error(`ball: dimension must be a positive integer, got ${d}`);
+	if (p !== 2) throw new Error("ball: only the Euclidean p=2 case is supported");
+	const [k1, k2] = split(key$1, 2);
+	const z = normal(k1, [...shape$1, d]);
+	const norm = sqrt(z.ref.mul(z.ref).sum(-1, { keepdims: true }));
+	const radius = exp(log(uniform(k2, [...shape$1, 1])).mul(1 / d));
+	return z.div(norm).mul(radius);
+}, { staticArgnums: [1, 2] });
+/**
 * Sample Bernoulli random variables with given mean (0,1 categorical).
 *
 * Returns a random Boolean array with the specified shape. `p` can be an array
@@ -8025,6 +8200,57 @@ const cauchy = jit$1(function cauchy$1(key$1, shape$1 = []) {
 	return tan(u.sub(.5).mul(Math.PI));
 }, { staticArgnums: [1] });
 /**
+* Sample from a population with optional replacement and optional probabilities.
+*
+* This implements the common JAX-compatible cases: integer populations and
+* array populations along `axis`. Probabilities `p`, if provided, are sampled
+* via `categorical(log(p))`.
+*/
+function choice(key$1, a, { shape: shape$1 = [], replace = true, p, axis = 0 } = {}) {
+	let n;
+	let values = null;
+	if (typeof a === "number") {
+		if (!Number.isInteger(a) || a < 0) throw new Error(`choice: a must be a non-negative integer, got ${a}`);
+		n = a;
+	} else {
+		values = fudgeArray(a);
+		axis = checkAxis(axis, values.ndim);
+		n = values.shape[axis];
+	}
+	let indices;
+	if (p !== void 0) indices = categorical(key$1, log(p), {
+		shape: shape$1,
+		replace
+	});
+	else if (replace) indices = randint(key$1, {
+		minval: 0,
+		maxval: n,
+		shape: shape$1
+	});
+	else {
+		const k = shape$1.reduce((acc, x) => acc * x, 1);
+		if (k > n) throw new Error(`Number of samples without replacement (${k}) cannot exceed population size (${n}).`);
+		indices = permutation(key$1, n).slice([0, k]).reshape(shape$1);
+	}
+	if (values === null) return indices;
+	const index = JsArray(axis).fill([]);
+	index.push(indices);
+	return values.slice(...index);
+}
+/**
+* @function
+* Sample double-sided Maxwell random values with the provided location and scale.
+*/
+const doubleSidedMaxwell = jit$1(function doubleSidedMaxwell$1(key$1, loc, scale, shape$1 = []) {
+	loc = fudgeArray(loc);
+	scale = fudgeArray(scale);
+	const [k1, k2] = split(key$1, 2);
+	return rademacher(k1, {
+		shape: shape$1,
+		dtype: DType.Float32
+	}).mul(maxwell(k2, shape$1)).mul(scale).add(loc);
+}, { staticArgnums: [3] });
+/**
 * @function
 * Sample exponential random values according to `p(x) = exp(-x)`.
 */
@@ -8034,6 +8260,14 @@ const exponential = jit$1(function exponential$1(key$1, shape$1 = []) {
 }, { staticArgnums: [1] });
 /**
 * @function
+* Sample geometric random values: the number of trials until first success.
+*/
+const geometric = jit$1(function geometric$1(key$1, p, { shape: shape$1 = [], dtype = DType.Int32 } = {}) {
+	p = fudgeArray(p);
+	return floor(log1p(negative(uniform(key$1, shape$1))).div(log1p(negative(p)))).add(1).astype(dtype);
+}, { staticArgnums: [2] });
+/**
+* @function
 * Sample from a Gumbel distribution with location 0 and scale 1.
 *
 * Uses inverse transform sampling: `x = -log(-log(u))` where u ~ Uniform(0, 1).
@@ -8058,6 +8292,32 @@ const laplace = jit$1(function laplace$1(key$1, shape$1 = []) {
 }, { staticArgnums: [1] });
 /**
 * @function
+* Sample from a logistic distribution with location 0 and scale 1.
+*
+* Uses inverse transform sampling: `x = log(u) - log(1-u)`.
+*/
+const logistic = jit$1(function logistic$1(key$1, shape$1 = []) {
+	const u = uniform(key$1, shape$1);
+	return log(u.ref).sub(log1p(negative(u)));
+}, { staticArgnums: [1] });
+/**
+* @function
+* Sample log-normal random values: `exp(sigma * normal(key, shape))`.
+*/
+const lognormal = jit$1(function lognormal$1(key$1, sigma = 1, shape$1 = []) {
+	sigma = fudgeArray(sigma);
+	return exp(normal(key$1, shape$1).mul(sigma));
+}, { staticArgnums: [2] });
+/**
+* @function
+* Sample Maxwell-distributed random values.
+*/
+const maxwell = jit$1(function maxwell$1(key$1, shape$1 = []) {
+	const z = normal(key$1, [...shape$1, 3]);
+	return sqrt(z.ref.mul(z).sum(-1));
+}, { staticArgnums: [1] });
+/**
+* @function
 * Sample multivariate normal random values with given mean and covariance.
 *
 * The values are returned with the given shape, along with the final dimension
@@ -8098,6 +8358,97 @@ const normal = jit$1(function normal$1(key$1, shape$1 = []) {
 	const theta = u2.mul(2 * Math.PI);
 	return radius.mul(cos(theta));
 }, { staticArgnums: [1] });
+/**
+* @function
+* Sample from a Pareto distribution with shape parameter `b` and support [1, ∞).
+*/
+const pareto = jit$1(function pareto$1(key$1, b, shape$1 = []) {
+	b = fudgeArray(b);
+	return exp(exponential(key$1, shape$1).div(b));
+}, { staticArgnums: [2] });
+/**
+* Return a random permutation of an integer range or of an array along `axis`.
+*/
+function permutation(key$1, x, axis = 0) {
+	if (typeof x === "number") {
+		if (!Number.isInteger(x) || x < 0) throw new Error(`permutation: x must be a non-negative integer, got ${x}`);
+		return argsort(uniform(key$1, [x])).astype(DType.Int32);
+	}
+	const arr = fudgeArray(x);
+	axis = checkAxis(axis, arr.ndim);
+	const perm = permutation(key$1, arr.shape[axis]);
+	const index = JsArray(axis).fill([]);
+	index.push(perm);
+	return arr.slice(...index);
+}
+/**
+* @function
+* Sample Rademacher random values, uniformly from {-1, 1}.
+*/
+const rademacher = jit$1(function rademacher$1(key$1, { shape: shape$1 = [], dtype = DType.Int32 } = {}) {
+	if (dtype === DType.Uint32 || dtype === DType.Bool) throw new Error(`rademacher: unsupported dtype ${dtype}`);
+	const one = array(1, {
+		dtype,
+		device: key$1.device
+	});
+	const minusOne = array(-1, {
+		dtype,
+		device: key$1.device
+	});
+	return where(bernoulli(key$1, .5, shape$1), one, minusOne);
+}, { staticArgnums: [1] });
+/**
+* @function
+* Sample integer values uniformly from `[minval, maxval)`.
+*
+* This uses modulo reduction of uniform 32-bit random bits. For ranges that do
+* not divide 2^32, this introduces a very small modulo bias.
+*/
+const randint = jit$1(function randint$1(key$1, { minval, maxval, shape: shape$1 = [], dtype = DType.Int32 }) {
+	if (!Number.isInteger(minval) || !Number.isInteger(maxval)) throw new Error("randint: minval and maxval must be integers");
+	if (minval >= maxval) throw new Error(`Invalid range: [${minval}, ${maxval}).`);
+	if (dtype !== DType.Int32 && dtype !== DType.Uint32) throw new Error(`randint: dtype must be int32 or uint32, got ${dtype}`);
+	if (dtype === DType.Uint32 && minval < 0) throw new Error("randint: uint32 dtype requires minval >= 0");
+	const range$1 = maxval - minval;
+	return bits(key$1, shape$1).mod(range$1).astype(dtype).add(minval);
+}, { staticArgnums: [1] });
+/**
+* @function
+* Sample Rayleigh random values with the provided scale parameter.
+*/
+const rayleigh = jit$1(function rayleigh$1(key$1, scale = 1, shape$1 = []) {
+	scale = fudgeArray(scale);
+	return sqrt(exponential(key$1, shape$1).mul(2)).mul(scale);
+}, { staticArgnums: [2] });
+/**
+* @function
+* Sample triangular random values on `[left, right]` with the given mode.
+*/
+const triangular = jit$1(function triangular$1(key$1, left, mode, right, shape$1 = []) {
+	left = fudgeArray(left);
+	mode = fudgeArray(mode);
+	right = fudgeArray(right);
+	const u = uniform(key$1, shape$1);
+	const width = right.ref.sub(left.ref);
+	const leftSpan = mode.ref.sub(left.ref);
+	const rightSpan = right.ref.sub(mode);
+	const cutoff = leftSpan.ref.div(width.ref);
+	const cond = u.ref.less(cutoff);
+	const lower = left.add(sqrt(u.ref.mul(width.ref).mul(leftSpan)));
+	const upper = right.sub(sqrt(negative(u).add(1).mul(width).mul(rightSpan)));
+	return where(cond, lower, upper);
+}, { staticArgnums: [4] });
+/**
+* @function
+* Sample Weibull minimum random values.
+*
+* Uses `scale * exponential(key) ** (1 / concentration)`.
+*/
+const weibullMin = jit$1(function weibullMin$1(key$1, scale, concentration, shape$1 = []) {
+	scale = fudgeArray(scale);
+	concentration = fudgeArray(concentration);
+	return scale.mul(exp(log(exponential(key$1, shape$1)).div(concentration)));
+}, { staticArgnums: [3] });
 
 //#endregion
 //#region src/library/scipy-special.ts
@@ -8291,4 +8642,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-BhsnpeB0.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-DI-V78Rk.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-CD3WK_Me.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-DMauYnfl.cjs');
+const require_backend = require('./backend-x-6vqzIM.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]);