@jax-js/jax 0.1.6 → 0.1.7

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, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-Dx6Ob2D1.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, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-nEolvdLv.js";
 
 //#region src/frontend/convolution.ts
 /**
@@ -306,11 +306,11 @@ function map(fn, tree, ...rest) {
 }
 /** Take a reference of every array in a tree. */
 function ref(tree) {
-	return map((x) => x.ref, tree);
+	return map((x) => x instanceof Tracer ? x.ref : x, tree);
 }
 /** Dispose every array in a tree. */
 function dispose(tree) {
-	if (tree) map((x) => x.dispose(), tree);
+	if (tree) map((x) => x instanceof Tracer ? x.dispose() : void 0, tree);
 }
 
 //#endregion
@@ -584,14 +584,20 @@ function shrink(x, slice) {
 }
 function pad$1(x, width) {
 	const nd = ndim$1(x);
-	if (typeof width === "number") width = [[width, width]];
-	else if (isNumberPair(width)) width = [width];
-	else if (!Array.isArray(width) || !width.every(isNumberPair)) throw new TypeError(`Invalid pad() type: ${JSON.stringify(width)}`);
-	if (width.length === 1) {
-		const [w0, w1] = width[0];
-		width = rep(nd, () => [w0, w1]);
-	} else if (width.length !== nd) throw new Error(`Invalid pad(): expected ${nd} axes, got ${width.length}`);
-	return bind1(Primitive.Pad, [x], { width });
+	let w;
+	if (typeof width === "number") w = [[width, width]];
+	else if (isNumberPair(width)) w = [width];
+	else if (!Array.isArray(width)) {
+		const indicesAndPairs = Object.entries(width);
+		w = rep(nd, [0, 0]);
+		for (const [k, v] of indicesAndPairs) w[checkAxis(parseInt(k), nd)] = v;
+	} else if (!width.every(isNumberPair)) throw new TypeError(`Invalid pad() type: ${JSON.stringify(width)}`);
+	else w = width;
+	if (w.length === 1) {
+		const [w0, w1] = w[0];
+		w = rep(nd, () => [w0, w1]);
+	} else if (w.length !== nd) throw new Error(`Invalid pad(): expected ${nd} axes, got ${w.length}`);
+	return bind1(Primitive.Pad, [x], { width: w });
 }
 function triangularSolve$1(a, b, { lower = false, unitDiagonal = false } = {}) {
 	const as = getShape(a);
@@ -767,6 +773,22 @@ var Tracer = class Tracer {
 		const result = reduce(this.astype(castDtype), AluOp.Add, axis, opts);
 		return result.mul(1 / n).astype(originalDtype);
 	}
+	/** Minimum of the elements of the array along a given axis. */
+	min(axis = null, opts) {
+		return reduce(this, AluOp.Min, axis, opts);
+	}
+	/** Maximum of the elements of the array along a given axis. */
+	max(axis = null, opts) {
+		return reduce(this, AluOp.Max, axis, opts);
+	}
+	/** Test whether all array elements along a given axis evaluate to true. */
+	all(axis = null, opts) {
+		return this.astype(DType.Bool).min(axis, opts);
+	}
+	/** Test whether any array element along a given axis evaluates to true. */
+	any(axis = null, opts) {
+		return this.astype(DType.Bool).max(axis, opts);
+	}
 	/** Permute the dimensions of an array. Defaults to reversing the axis order. */
 	transpose(perm) {
 		return transpose$1(this, perm);
@@ -5570,6 +5592,7 @@ __export(numpy_exports, {
 	moveaxis: () => moveaxis$1,
 	multiply: () => multiply,
 	nan: () => nan,
+	nanToNum: () => nanToNum,
 	ndim: () => ndim,
 	negative: () => negative,
 	notEqual: () => notEqual,
@@ -5767,24 +5790,22 @@ function max(a, axis = null, opts) {
 	return reduce(a, AluOp.Max, axis, opts);
 }
 /**
-* Test whether all array elements along a given axis evaluate to True.
+* Test whether any array element along a given axis evaluates to True.
 *
 * Returns a boolean array with the same shape as `a` with the specified axis
 * removed. If axis is None, returns a scalar.
 */
-function all(a, axis = null, opts) {
-	a = fudgeArray(a).astype(DType.Bool);
-	return min(a, axis, opts);
+function any(a, axis = null, opts) {
+	return fudgeArray(a).any(axis, opts);
 }
 /**
-* Test whether any array element along a given axis evaluates to True.
+* Test whether all array elements along a given axis evaluate to True.
 *
 * Returns a boolean array with the same shape as `a` with the specified axis
 * removed. If axis is None, returns a scalar.
 */
-function any(a, axis = null, opts) {
-	a = fudgeArray(a).astype(DType.Bool);
-	return max(a, axis, opts);
+function all(a, axis = null, opts) {
+	return fudgeArray(a).all(axis, opts);
 }
 /** Return the peak-to-peak range along a given axis (`max - min`). */
 function ptp(a, axis = null, opts) {
@@ -5885,7 +5906,7 @@ function split$1(a, indicesOrSections, axis = 0) {
 		const partSize = size$1 / indicesOrSections;
 		sizes = rep(indicesOrSections, partSize);
 	} else {
-		const indices = indicesOrSections;
+		const indices = indicesOrSections.map((i) => i < 0 ? i + size$1 : i);
 		sizes = [indices[0]];
 		for (let i = 1; i < indices.length; i++) sizes.push(indices[i] - indices[i - 1]);
 		sizes.push(size$1 - indices[indices.length - 1]);
@@ -6833,6 +6854,21 @@ function isposinf(x) {
 	return isFloatDtype(x.dtype) ? x.equal(Infinity) : fullLike$1(x, false);
 }
 /**
+* Replace NaN and infinite entries in an array.
+*
+* By default, NaNs are replaced with `0.0`, and infinities are are substituted
+* with the corresponding maximum or minimum finite values.
+*/
+function nanToNum(x, { nan: nan$1 = 0, posinf = null, neginf = null } = {}) {
+	x = fudgeArray(x);
+	x = where(isnan(x.ref), nan$1, x);
+	posinf ??= isFloatDtype(x.dtype) ? finfo(x.dtype).max : iinfo(x.dtype).max;
+	neginf ??= isFloatDtype(x.dtype) ? finfo(x.dtype).min : iinfo(x.dtype).min;
+	x = where(isposinf(x.ref), posinf, x);
+	x = where(isneginf(x.ref), neginf, x);
+	return x;
+}
+/**
 * @function
 * Test element-wise for finite values (not infinity or NaN).
 */
@@ -7541,8 +7577,6 @@ function oneHot(x, numClasses) {
 *   `[B, L, N, H]`, or `[L, N, H]` if `B` is omitted.
 */
 function dotProductAttention(query, key$1, value, opts = {}) {
-	if (opts.querySeqLengths !== void 0 || opts.keyValueSeqLengths !== void 0) throw new Error("Sequence length masking is not yet implemented");
-	if (opts.localWindowSize !== void 0) throw new Error("Local attention is not yet implemented");
 	query = fudgeArray(query);
 	key$1 = fudgeArray(key$1);
 	value = fudgeArray(value);
@@ -7580,6 +7614,38 @@ function dotProductAttention(query, key$1, value, opts = {}) {
 		const causalMask = tri(L, S, 0, { dtype: DType.Bool });
 		scores = where(causalMask, scores, -Infinity);
 	}
+	if (opts.localWindowSize !== void 0) {
+		const [before, after] = typeof opts.localWindowSize === "number" ? [opts.localWindowSize, opts.localWindowSize] : opts.localWindowSize;
+		if (before < 0 || after < 0 || !Number.isInteger(before) || !Number.isInteger(after)) throw new Error(`dotProductAttention: localWindowSize values must be non-negative, got ${opts.localWindowSize}`);
+		const localMask = tri(L, S, after, { dtype: DType.Bool }).mul(tri(L, S, -before - 1, { dtype: DType.Bool }).notEqual(true));
+		scores = where(localMask, scores, -Infinity);
+	}
+	if (opts.querySeqLengths !== void 0) {
+		const sl = expandDims(opts.querySeqLengths, [
+			-1,
+			-2,
+			-3
+		]);
+		scores = where(arange(L).reshape([
+			1,
+			1,
+			L,
+			1
+		]).less(sl), scores, -Infinity);
+	}
+	if (opts.keyValueSeqLengths !== void 0) {
+		const sl = expandDims(opts.keyValueSeqLengths, [
+			-1,
+			-2,
+			-3
+		]);
+		scores = where(arange(S).reshape([
+			1,
+			1,
+			1,
+			S
+		]).less(sl), scores, -Infinity);
+	}
 	const attn = softmax(scores, -1);
 	const out = einsum("BNLS,BSNH->BLNH", attn, value);
 	return isRank3 ? out.reshape([

package/dist/{webgl-CyfzNW8T.cjs → webgl-DIIbKJ0G.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-D7s-Retx.cjs');
+const require_backend = require('./backend-B3foXiV_.cjs');
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `

package/dist/{webgl-CLLvzJlO.js → webgl-DweKSWEm.js}

@@ -1,4 +1,4 @@
-import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-Dx6Ob2D1.js";
+import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-nEolvdLv.js";
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `

package/dist/{webgpu-C-VfevQW.js → webgpu-B96vzWGE.js}

@@ -1,4 +1,4 @@
-import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, findPow2, isFloatDtype, mapSetUnion, prod, range, strip1, tuneWebgpu } from "./backend-Dx6Ob2D1.js";
+import { AluExp, AluGroup, AluOp, DEBUG, DType, Executable, FpHash, Routines, SlotError, UnsupportedOpError, UnsupportedRoutineError, findPow2, isFloatDtype, mapSetUnion, prod, range, strip1, tuneWebgpu } from "./backend-nEolvdLv.js";
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `

package/dist/{webgpu-rraa6dfz.cjs → webgpu-BykvF26B.cjs}

@@ -1,4 +1,4 @@
-const require_backend = require('./backend-D7s-Retx.cjs');
+const require_backend = require('./backend-B3foXiV_.cjs');
 
 //#region src/backend/webgpu/builtins.ts
 const threefrySrc = `

package/package.json

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