@jax-js/jax 0.1.11 → 0.1.13

package/README.md

@@ -63,6 +63,7 @@ of late 2025.
 
 Community usage:
 
+- [**g9-jaxjs**: Automatically interactive graphics with forward-mode AD](https://srush.github.io/g9jax/)
 - [**autoresearch-webgpu**: autoresesarch, in the browser](https://autoresearch.lucasgelfond.online/)
 - [**tanh.xyz**: Interactive ML visualizations](https://tanh.xyz/)
 - [**jax-js-bayes**: Declarative Bayesian modeling library](https://github.com/StefanSko/jax-js-bayes)
@@ -72,10 +73,13 @@ Demos on the jax-js website:
 - [Training neural networks on MNIST](https://jax-js.com/mnist)
 - [Voice cloning: Kyutai Pocket TTS](https://jax-js.com/tts)
 - [CLIP embeddings for books in-browser](https://jax-js.com/mobileclip)
+- [Object detection: D-FINE (ONNX)](https://jax-js.com/d-fine)
 - [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
 - [Fluid simulation (Navier-Stokes)](https://jax-js.com/fluid-sim)
+- [Neural cellular automata](https://jax-js.com/nca-growing)
 - [In-browser REPL](https://jax-js.com/repl)
 - [Matmul benchmark](https://jax-js.com/bench/matmul)
+- [Matvec benchmark](https://jax-js.com/bench/matvec)
 - [Conv2d benchmark](https://jax-js.com/bench/conv2d)
 - [Mandelbrot set](https://jax-js.com/mandelbrot)
 
@@ -422,7 +426,6 @@ Contributions are welcomed! Some fruitful areas to look into:
 - Adding support for more JAX functions and operations, see [compatibility table](./FEATURES.md).
 - Improving performance of the WebGPU and Wasm runtimes, generating better kernels, and using SIMD
   and multithreading. (Even single-threaded Wasm could be ~20x faster.)
-- Helping the JIT compiler to fuse operations in more cases, like `tanh` branches.
 - Making a fast transformer inference engine, comparing against onnxruntime-web.
 
 You may join our [Discord server](https://discord.gg/BW6YsCd4Tf) and chat with the community.

package/dist/{backend-DZvR7mZV.js → backend-DLEk-B3V.js}

@@ -1430,11 +1430,13 @@ var Reduction = class {
 function accessorGlobal(dtype, gid, st, indices) {
 	const [index, valid] = st.toAluExp(indices);
 	const [, len] = st.views[0].dataRange();
+	if (valid.resolve()) return AluExp.globalIndex(dtype, gid, len, index);
 	return AluExp.where(valid, AluExp.globalIndex(dtype, gid, len, index), AluExp.const(dtype, 0));
 }
 /** Expression for accessing `indices` in an array recipe with variable "idx". */
 function accessorAluExp(exp, st, indices) {
 	const [index, valid] = st.toAluExp(indices);
+	if (valid.resolve()) return exp.substitute({ idx: index });
 	return AluExp.where(valid, exp.substitute({ idx: index }), AluExp.const(exp.dtype, 0));
 }
 function threefry2x32(k0, k1, c0, c1) {
@@ -4520,7 +4522,7 @@ var WasmBackend = class {
 		const buffer = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
 		if (count === void 0) count = buffer.byteLength - start;
-		if (buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
+		if (hasSharedArrayBuffer() && buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
 		else return buffer.slice(start, start + count);
 	}
 	async prepareKernel(kernel) {
@@ -5059,7 +5061,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-Dg8FpYrH.js");
+		const { WebGPUBackend } = await import("./webgpu-NkF1TZ0t.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -5097,7 +5099,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await import("./webgl-D8-14NzA.js");
+		const { WebGLBackend } = await import("./webgl-NsFtyIts.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

package/dist/{backend-DlYlOYqN.cjs → backend-DMyuoWi2.cjs}

@@ -1431,11 +1431,13 @@ var Reduction = class {
 function accessorGlobal(dtype, gid, st, indices) {
 	const [index, valid] = st.toAluExp(indices);
 	const [, len] = st.views[0].dataRange();
+	if (valid.resolve()) return AluExp.globalIndex(dtype, gid, len, index);
 	return AluExp.where(valid, AluExp.globalIndex(dtype, gid, len, index), AluExp.const(dtype, 0));
 }
 /** Expression for accessing `indices` in an array recipe with variable "idx". */
 function accessorAluExp(exp, st, indices) {
 	const [index, valid] = st.toAluExp(indices);
+	if (valid.resolve()) return exp.substitute({ idx: index });
 	return AluExp.where(valid, exp.substitute({ idx: index }), AluExp.const(exp.dtype, 0));
 }
 function threefry2x32(k0, k1, c0, c1) {
@@ -4521,7 +4523,7 @@ var WasmBackend = class {
 		const buffer = this.#getBuffer(slot);
 		if (start === void 0) start = 0;
 		if (count === void 0) count = buffer.byteLength - start;
-		if (buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
+		if (hasSharedArrayBuffer() && buffer.buffer instanceof SharedArrayBuffer) return new Uint8Array(buffer.slice(start, start + count));
 		else return buffer.slice(start, start + count);
 	}
 	async prepareKernel(kernel) {
@@ -5060,7 +5062,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-uU9nnttc.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-DDGCYtHa.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -5098,7 +5100,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-Ovaaa-Qx.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-pbfUGDA6.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

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-DlYlOYqN.cjs');
+const require_backend = require('./backend-DMyuoWi2.cjs');
 
 //#region src/frontend/convolution.ts
 /**
@@ -240,7 +240,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";
@@ -288,7 +288,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) {
@@ -2495,7 +2495,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;
@@ -2935,6 +2935,15 @@ var Array$1 = class Array$1 extends Tracer {
 		this.#check();
 		const indices = require_backend.unravelAlu(this.#st.shape, require_backend.AluVar.gidx);
 		if (this.#source instanceof require_backend.AluExp) {
+			let resolvedSource;
+			if (this.#st.contiguous && this.#st.size < inlineArrayLimit && (resolvedSource = this.#source.resolve()) !== void 0) {
+				const byteLength = this.#st.size * require_backend.byteWidth(this.#dtype);
+				const initialData = new Uint8Array(byteLength);
+				require_backend.dtypedArray(this.#dtype, initialData).fill(resolvedSource);
+				this.#source = this.#backend.malloc(byteLength, initialData);
+				this.#st = require_backend.ShapeTracker.fromShape(this.shape);
+				return;
+			}
 			const exp$2 = require_backend.accessorAluExp(this.#source, this.#st, indices);
 			const kernel = new require_backend.Kernel(0, this.#st.size, exp$2);
 			const output = this.#backend.malloc(kernel.bytes);
@@ -3385,7 +3394,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];
 			}
@@ -4269,7 +4278,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]);
@@ -4281,7 +4290,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,
@@ -4293,11 +4302,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 }) {
@@ -5379,8 +5388,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);
@@ -5469,8 +5478,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);
@@ -5508,9 +5517,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,
@@ -7366,7 +7375,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);
@@ -7553,7 +7562,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];
@@ -8072,19 +8081,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.`);
@@ -8137,6 +8161,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
@@ -8198,6 +8237,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 = require_backend.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: require_backend.DType.Float32
+	}).mul(maxwell(k2, shape$1)).mul(scale).add(loc);
+}, { staticArgnums: [3] });
+/**
 * @function
 * Sample exponential random values according to `p(x) = exp(-x)`.
 */
@@ -8207,6 +8297,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 = require_backend.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).
@@ -8231,6 +8329,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
@@ -8271,6 +8395,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(require_backend.DType.Int32);
+	}
+	const arr = fudgeArray(x);
+	axis = require_backend.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 = require_backend.DType.Int32 } = {}) {
+	if (dtype === require_backend.DType.Uint32 || dtype === require_backend.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 = require_backend.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 !== require_backend.DType.Int32 && dtype !== require_backend.DType.Uint32) throw new Error(`randint: dtype must be int32 or uint32, got ${dtype}`);
+	if (dtype === require_backend.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

package/dist/index.d.cts

@@ -2722,7 +2722,7 @@ declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: Ar
   localWindowSize?: number | [number, number];
 }): Array;
 declare namespace random_d_exports {
-  export { bernoulli, bits, categorical, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
+  export { ball, bernoulli, bits, categorical, cauchy, choice, doubleSidedMaxwell, exponential, geometric, gumbel, key, laplace, logistic, lognormal, maxwell, multivariateNormal, normal, pareto, permutation, rademacher, randint, rayleigh, split, triangular, uniform, weibullMin };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: ArrayLike): Array;
@@ -2738,6 +2738,16 @@ declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefin
   minval?: number | undefined;
   maxval?: number | undefined;
 } | undefined) => Array>;
+/**
+ * @function
+ * Sample points uniformly from the Euclidean unit ball in `d` dimensions.
+ *
+ * Only the Euclidean `p=2` case is currently supported.
+ */
+declare const ball: OwnedFunction<(key: ArrayLike, d: number, args_2?: {
+  p?: number | undefined;
+  shape?: number[] | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -2778,11 +2788,42 @@ declare const categorical: OwnedFunction<(key: ArrayLike, logits: ArrayLike, arg
  * Uses inverse transform sampling: `x = tan(π * (u - 0.5))` where u ~ Uniform(0, 1).
  */
 declare const cauchy: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * 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))`.
+ */
+declare function choice(key: Array, a: number | ArrayLike, {
+  shape,
+  replace,
+  p,
+  axis
+}?: {
+  shape?: number[];
+  replace?: boolean;
+  p?: ArrayLike;
+  axis?: number;
+}): Array;
+/**
+ * @function
+ * Sample double-sided Maxwell random values with the provided location and scale.
+ */
+declare const doubleSidedMaxwell: OwnedFunction<(key: ArrayLike, loc: ArrayLike, scale: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample exponential random values according to `p(x) = exp(-x)`.
  */
 declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample geometric random values: the number of trials until first success.
+ */
+declare const geometric: OwnedFunction<(key: ArrayLike, p: ArrayLike, args_2?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
 /**
  * @function
  * Sample from a Gumbel distribution with location 0 and scale 1.
@@ -2798,6 +2839,23 @@ declare const gumbel: OwnedFunction<(key: ArrayLike, shape?: number[] | undefine
  * Inverting: `x = -sign(u - 0.5) * log(1 - 2 * |u - 0.5|)`.
  */
 declare const laplace: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a logistic distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: `x = log(u) - log(1-u)`.
+ */
+declare const logistic: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample log-normal random values: `exp(sigma * normal(key, shape))`.
+ */
+declare const lognormal: OwnedFunction<(key: ArrayLike, sigma?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Maxwell-distributed random values.
+ */
+declare const maxwell: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample multivariate normal random values with given mean and covariance.
@@ -2824,6 +2882,53 @@ declare const multivariateNormal: OwnedFunction<(key: ArrayLike, mean: ArrayLike
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a Pareto distribution with shape parameter `b` and support [1, ∞).
+ */
+declare const pareto: OwnedFunction<(key: ArrayLike, b: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * Return a random permutation of an integer range or of an array along `axis`.
+ */
+declare function permutation(key: Array, x: number | ArrayLike, axis?: number): Array;
+/**
+ * @function
+ * Sample Rademacher random values, uniformly from {-1, 1}.
+ */
+declare const rademacher: OwnedFunction<(key: ArrayLike, args_1?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
+/**
+ * @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.
+ */
+declare const randint: OwnedFunction<(key: ArrayLike, args_1: {
+  minval: number;
+  maxval: number;
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+}) => Array>;
+/**
+ * @function
+ * Sample Rayleigh random values with the provided scale parameter.
+ */
+declare const rayleigh: OwnedFunction<(key: ArrayLike, scale?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample triangular random values on `[left, right]` with the given mode.
+ */
+declare const triangular: OwnedFunction<(key: ArrayLike, left: ArrayLike, mode: ArrayLike, right: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Weibull minimum random values.
+ *
+ * Uses `scale * exponential(key) ** (1 / concentration)`.
+ */
+declare const weibullMin: OwnedFunction<(key: ArrayLike, scale: ArrayLike, concentration: ArrayLike, shape?: number[] | undefined) => Array>;
 declare namespace scipy_special_d_exports {
   export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
 }