@jax-js/jax 0.1.2 → 0.1.3

package/dist/index.d.cts

@@ -254,7 +254,7 @@ declare class AluExp implements FpHashable {
   /** Substitute variables in this AluExp to values. */
   substitute(variables: Record<string, AluExp>): AluExp;
   /** Reindex gid values in this expression as needed. */
-  reindexGids(gidMap: Map<number, number>): AluExp;
+  reindexGids(newGids: number[]): AluExp;
   get min(): number;
   get max(): number;
   /** Largest known integer that divides self. */
@@ -464,7 +464,7 @@ declare class Executable<T = any> {
   data: T);
 }
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, cos, cosh, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, squeeze, stack, std, subtract, sum, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, cos, cosh, cumsum, cumulativeSum, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, squeeze, stack, std, subtract, sum, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -608,6 +608,15 @@ declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  * specified axis.
  */
 declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+/**
+ * Cumulative sum of elements along an axis.
+ *
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
+ */
+declare function cumsum(a: ArrayLike, axis?: number): Array;
+/** @function Alternative name for `jax.numpy.cumsum()`. */
+declare const cumulativeSum: typeof cumsum;
 /** Reverse the elements in an array along the given axes. */
 declare function flip(x: ArrayLike, axis?: Axis): Array;
 /**
@@ -1076,6 +1085,7 @@ declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefin
 //#region src/frontend/convolution.d.ts
 /** Definition of a general dilated convolution. Should be valid on creation. */
 interface ConvParams {
+  vmapDims: number;
   strides: number[];
   padding: [number, number][];
   lhsDilation: number[];
@@ -1083,6 +1093,10 @@ interface ConvParams {
 }
 /**
  * Check that the shapes and parameters passed to convolution are valid.
+ * Expected shapes of the lhs and rhs of the convolution are:
+ *
+ * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
+ * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
  *
  * If the check succeeds, returns the output shape.
  */
@@ -1745,10 +1759,12 @@ type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
  */
 declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
   lhsDilation,
-  rhsDilation
+  rhsDilation,
+  featureGroupCount
 }?: {
   lhsDilation?: number[];
   rhsDilation?: number[];
+  featureGroupCount?: number;
 }): Array;
 /** Convenience wrapper around `convGeneralDilated`. */
 declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
@@ -1915,9 +1931,9 @@ declare function logSoftmax(x: ArrayLike, axis?: Axis): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/LogSumExp
  */
-declare function logsumexp(x: ArrayLike, axis?: Axis): Array;
+declare function logsumexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
-declare function logmeanexp(x: ArrayLike, axis?: Axis): Array;
+declare function logmeanexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
  * Standardizes input to zero mean and unit variance.
  *

package/dist/index.d.ts

@@ -251,7 +251,7 @@ declare class AluExp implements FpHashable {
   /** Substitute variables in this AluExp to values. */
   substitute(variables: Record<string, AluExp>): AluExp;
   /** Reindex gid values in this expression as needed. */
-  reindexGids(gidMap: Map<number, number>): AluExp;
+  reindexGids(newGids: number[]): AluExp;
   get min(): number;
   get max(): number;
   /** Largest known integer that divides self. */
@@ -461,7 +461,7 @@ declare class Executable<T = any> {
   data: T);
 }
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, cos, cosh, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, squeeze, stack, std, subtract, sum, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, cos, cosh, cumsum, cumulativeSum, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, squeeze, stack, std, subtract, sum, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -605,6 +605,15 @@ declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  * specified axis.
  */
 declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+/**
+ * Cumulative sum of elements along an axis.
+ *
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
+ */
+declare function cumsum(a: ArrayLike, axis?: number): Array;
+/** @function Alternative name for `jax.numpy.cumsum()`. */
+declare const cumulativeSum: typeof cumsum;
 /** Reverse the elements in an array along the given axes. */
 declare function flip(x: ArrayLike, axis?: Axis): Array;
 /**
@@ -1073,6 +1082,7 @@ declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefin
 //#region src/frontend/convolution.d.ts
 /** Definition of a general dilated convolution. Should be valid on creation. */
 interface ConvParams {
+  vmapDims: number;
   strides: number[];
   padding: [number, number][];
   lhsDilation: number[];
@@ -1080,6 +1090,10 @@ interface ConvParams {
 }
 /**
  * Check that the shapes and parameters passed to convolution are valid.
+ * Expected shapes of the lhs and rhs of the convolution are:
+ *
+ * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
+ * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
  *
  * If the check succeeds, returns the output shape.
  */
@@ -1742,10 +1756,12 @@ type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
  */
 declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
   lhsDilation,
-  rhsDilation
+  rhsDilation,
+  featureGroupCount
 }?: {
   lhsDilation?: number[];
   rhsDilation?: number[];
+  featureGroupCount?: number;
 }): Array;
 /** Convenience wrapper around `convGeneralDilated`. */
 declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
@@ -1912,9 +1928,9 @@ declare function logSoftmax(x: ArrayLike, axis?: Axis): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/LogSumExp
  */
-declare function logsumexp(x: ArrayLike, axis?: Axis): Array;
+declare function logsumexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
-declare function logmeanexp(x: ArrayLike, axis?: Axis): Array;
+declare function logmeanexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
  * Standardizes input to zero mean and unit variance.
  *