@jax-js/jax 0.0.4 → 0.0.5

package/dist/index.d.cts

@@ -180,12 +180,12 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  * **Type lattice:**
  * ```text
  * bool -> uint32 -> int32 -> float16 -> float32
- *  weak f* --^
+ *  weakType --^
  * ```
  *
- * The asterisk f* is a weak type used for JS number constants. When creating
- * arrays, JS numbers default to float32 but "weak" so they cast to the dtype of
- * any array they are first combined with.
+ * `weakType` represents weakly typed arrays. These are created for JS numbers,
+ * which default to float32 but "weak" so they cast to the dtype of any array
+ * they are first combined with, except `bool`.
  *
  * **Examples:**
  * - `promoteTypes(bool, int32) → int32`
@@ -613,6 +613,7 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
     outDim: number;
   };
   [Primitive.JitCall]: {
+    name: string;
     jaxpr: Jaxpr;
     numConsts: number;
   };
@@ -651,10 +652,40 @@ declare abstract class Trace {
   abstract lift(val: Tracer): Tracer;
   abstract processPrimitive<P extends Primitive>(primitive: P, tracers: Tracer[], params: PrimitiveParams<P>): Tracer[];
 }
+/** Internal representation of an array value. */
 interface AbstractValue {
+  /** Shape of the array. Must be a static tuple of non-negative dimensions. */
   shape: number[];
+  /** Concrete data type of array elements. */
   dtype: DType;
+  /**
+   * Arrays created from JavaScript numbers (e.g., `np.array(3)`) are created as
+   * _weakly typed_ unless a dtype is explicitly specified.
+   *
+   * Weakly typed values will automatically cast to the data type of other
+   * arrays when used as an operand as an expression. This property only affects
+   * how they promote in type casting; their memory layout is still determined
+   * by the actual `dtype` field.
+   *
+   * ```ts
+   * const x = np.array(3); // weakType = true, dtype = float32
+   * const y = np.array([1, 2], { dtype: np.int32 }); // weakType = false, dtype = int32
+   * const z = x.add(y); // z has dtype int32 because x is weakly typed
+   * ```
+   *
+   * Weak types are present in JIT programs in their spec (e.g., Jaxpr inputs
+   * and outputs can be weakly typed) form. But they're solely a frontend
+   * concept. Backends are not aware of weak types.
+   */
+  weakType: boolean;
 }
+/**
+ * Broadcast shapes and promote types with casting for two avals.
+ *
+ * This implements the weak type behavior described in `promoteTypes()`, but not
+ * implemented in that function as `weakType` is not passed.
+ */
+
 declare abstract class Tracer {
   /** @ignore */
   readonly _trace: Trace;
@@ -712,8 +743,15 @@ declare abstract class Tracer {
   get shape(): number[];
   /** The total number of elements in the array. */
   get size(): number;
-  /** The dtype of the array. */
+  /** The dtype of elements stored in the array. */
   get dtype(): DType;
+  /**
+   * Whether the array is weakly typed.
+   *
+   * Weakly typed arrays will cast to the dtype of the other operand. See
+   * `promoteTypes()` for details.
+   */
+  get weakType(): boolean;
   /** The number of dimensions of the array. */
   get ndim(): number;
   /** @ignore */
@@ -805,7 +843,8 @@ declare abstract class Tracer {
 declare class ShapedArray implements AbstractValue {
   readonly shape: number[];
   readonly dtype: DType;
-  constructor(shape: number[], dtype: DType);
+  readonly weakType: boolean;
+  constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
   toString(): string;
@@ -841,6 +880,14 @@ type DTypeAndDevice = {
   dtype?: DType;
   device?: Device;
 };
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  dtype: DType;
+  weakType: boolean;
+  backend: Backend;
+  pending?: Iterable<PendingExecute>;
+};
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
  *
@@ -860,11 +907,7 @@ declare class Array extends Tracer {
    * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
    * will be freed when the array is disposed.
    */
-  constructor(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, {
-    pending
-  }?: {
-    pending?: Iterable<PendingExecute> | null;
-  });
+  constructor(args: ArrayConstructorArgs);
   /** @ignore */
   get aval(): ShapedArray;
   /** Return a simple string representation of the array's dimensions. */
@@ -926,8 +969,6 @@ declare class Array extends Tracer {
   static _implRules(): typeof implRules;
   _realizeSource(): number;
 }
-/** Construct an array from a single scalar constant. */
-
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
@@ -1480,10 +1521,10 @@ declare class Var {
 }
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
 declare class Lit {
-  readonly dtype: DType;
   readonly value: number;
   readonly aval: ShapedArray;
-  constructor(dtype: DType, value: number);
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
 }
 type Atom = Var | Lit;
 declare class VarPrinter {
@@ -1742,14 +1783,14 @@ declare function key(seed: number): Array;
 declare function split(key: Array, num?: number | number[]): Array;
 /** Sample uniform bits in the form of unsigned integers. */
 declare function bits(key: Array, shape?: number[]): Array;
-/** Sample uniform random values in [minval, maxval) with given shape. */
-declare function uniform(key: Array, shape?: number[], {
-  minval,
-  maxval
-}?: {
-  minval?: number;
-  maxval?: number;
-}): Array;
+/**
+ * @function
+ * Sample uniform random values in [minval, maxval) with given shape.
+ */
+declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined, args_2?: {
+  minval?: number | undefined;
+  maxval?: number | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -1757,16 +1798,20 @@ declare function uniform(key: Array, shape?: number[], {
  * and must be broadcastable to `shape`.
  */
 declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
-/** Sample exponential random values according to `p(x) = exp(-x)`. */
-declare function exponential(key: Array, shape?: number[]): Array;
 /**
+ * @function
+ * Sample exponential random values according to `p(x) = exp(-x)`.
+ */
+declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
  * Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
  *
  * Unlike JAX, this uses the Box-Muller transform. JAX uses the erf_inv primitive instead and
  * directly inverts the CDF, but we don't have support for that yet. Outputs will not be
  * bitwise identical to JAX.
  */
-declare function normal(key: Array, shape?: number[]): Array;
+declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 //#endregion
 //#region src/index.d.ts
 /**

package/dist/index.d.ts

@@ -177,12 +177,12 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  * **Type lattice:**
  * ```text
  * bool -> uint32 -> int32 -> float16 -> float32
- *  weak f* --^
+ *  weakType --^
  * ```
  *
- * The asterisk f* is a weak type used for JS number constants. When creating
- * arrays, JS numbers default to float32 but "weak" so they cast to the dtype of
- * any array they are first combined with.
+ * `weakType` represents weakly typed arrays. These are created for JS numbers,
+ * which default to float32 but "weak" so they cast to the dtype of any array
+ * they are first combined with, except `bool`.
  *
  * **Examples:**
  * - `promoteTypes(bool, int32) → int32`
@@ -610,6 +610,7 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
     outDim: number;
   };
   [Primitive.JitCall]: {
+    name: string;
     jaxpr: Jaxpr;
     numConsts: number;
   };
@@ -648,10 +649,40 @@ declare abstract class Trace {
   abstract lift(val: Tracer): Tracer;
   abstract processPrimitive<P extends Primitive>(primitive: P, tracers: Tracer[], params: PrimitiveParams<P>): Tracer[];
 }
+/** Internal representation of an array value. */
 interface AbstractValue {
+  /** Shape of the array. Must be a static tuple of non-negative dimensions. */
   shape: number[];
+  /** Concrete data type of array elements. */
   dtype: DType;
+  /**
+   * Arrays created from JavaScript numbers (e.g., `np.array(3)`) are created as
+   * _weakly typed_ unless a dtype is explicitly specified.
+   *
+   * Weakly typed values will automatically cast to the data type of other
+   * arrays when used as an operand as an expression. This property only affects
+   * how they promote in type casting; their memory layout is still determined
+   * by the actual `dtype` field.
+   *
+   * ```ts
+   * const x = np.array(3); // weakType = true, dtype = float32
+   * const y = np.array([1, 2], { dtype: np.int32 }); // weakType = false, dtype = int32
+   * const z = x.add(y); // z has dtype int32 because x is weakly typed
+   * ```
+   *
+   * Weak types are present in JIT programs in their spec (e.g., Jaxpr inputs
+   * and outputs can be weakly typed) form. But they're solely a frontend
+   * concept. Backends are not aware of weak types.
+   */
+  weakType: boolean;
 }
+/**
+ * Broadcast shapes and promote types with casting for two avals.
+ *
+ * This implements the weak type behavior described in `promoteTypes()`, but not
+ * implemented in that function as `weakType` is not passed.
+ */
+
 declare abstract class Tracer {
   /** @ignore */
   readonly _trace: Trace;
@@ -709,8 +740,15 @@ declare abstract class Tracer {
   get shape(): number[];
   /** The total number of elements in the array. */
   get size(): number;
-  /** The dtype of the array. */
+  /** The dtype of elements stored in the array. */
   get dtype(): DType;
+  /**
+   * Whether the array is weakly typed.
+   *
+   * Weakly typed arrays will cast to the dtype of the other operand. See
+   * `promoteTypes()` for details.
+   */
+  get weakType(): boolean;
   /** The number of dimensions of the array. */
   get ndim(): number;
   /** @ignore */
@@ -802,7 +840,8 @@ declare abstract class Tracer {
 declare class ShapedArray implements AbstractValue {
   readonly shape: number[];
   readonly dtype: DType;
-  constructor(shape: number[], dtype: DType);
+  readonly weakType: boolean;
+  constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
   toString(): string;
@@ -838,6 +877,14 @@ type DTypeAndDevice = {
   dtype?: DType;
   device?: Device;
 };
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  dtype: DType;
+  weakType: boolean;
+  backend: Backend;
+  pending?: Iterable<PendingExecute>;
+};
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
  *
@@ -857,11 +904,7 @@ declare class Array extends Tracer {
    * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
    * will be freed when the array is disposed.
    */
-  constructor(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, {
-    pending
-  }?: {
-    pending?: Iterable<PendingExecute> | null;
-  });
+  constructor(args: ArrayConstructorArgs);
   /** @ignore */
   get aval(): ShapedArray;
   /** Return a simple string representation of the array's dimensions. */
@@ -923,8 +966,6 @@ declare class Array extends Tracer {
   static _implRules(): typeof implRules;
   _realizeSource(): number;
 }
-/** Construct an array from a single scalar constant. */
-
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
@@ -1477,10 +1518,10 @@ declare class Var {
 }
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
 declare class Lit {
-  readonly dtype: DType;
   readonly value: number;
   readonly aval: ShapedArray;
-  constructor(dtype: DType, value: number);
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
 }
 type Atom = Var | Lit;
 declare class VarPrinter {
@@ -1739,14 +1780,14 @@ declare function key(seed: number): Array;
 declare function split(key: Array, num?: number | number[]): Array;
 /** Sample uniform bits in the form of unsigned integers. */
 declare function bits(key: Array, shape?: number[]): Array;
-/** Sample uniform random values in [minval, maxval) with given shape. */
-declare function uniform(key: Array, shape?: number[], {
-  minval,
-  maxval
-}?: {
-  minval?: number;
-  maxval?: number;
-}): Array;
+/**
+ * @function
+ * Sample uniform random values in [minval, maxval) with given shape.
+ */
+declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined, args_2?: {
+  minval?: number | undefined;
+  maxval?: number | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -1754,16 +1795,20 @@ declare function uniform(key: Array, shape?: number[], {
  * and must be broadcastable to `shape`.
  */
 declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
-/** Sample exponential random values according to `p(x) = exp(-x)`. */
-declare function exponential(key: Array, shape?: number[]): Array;
 /**
+ * @function
+ * Sample exponential random values according to `p(x) = exp(-x)`.
+ */
+declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
  * Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
  *
  * Unlike JAX, this uses the Box-Muller transform. JAX uses the erf_inv primitive instead and
  * directly inverts the CDF, but we don't have support for that yet. Outputs will not be
  * bitwise identical to JAX.
  */
-declare function normal(key: Array, shape?: number[]): Array;
+declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 //#endregion
 //#region src/index.d.ts
 /**