jax-js
    Preparing search index...

    Class Array

    A multidimensional numeric array with data stored on CPU or GPU.

    This is the library's core data type. Equivalent to jax.Array from JAX, or torch.Tensor.

    Not to be confused with the JavaScript "Array" constructor. Avoid importing this into your code's namespace if you're already using the JavaScript "Array" type by name.

    Hierarchy

    • Tracer
      • Array

    Properties

    id: number

    Accessors

    • get device(): Device

      Returns Device

    • get dtype(): DType

      The dtype of elements stored in the array.

      Returns DType

    • get ndim(): number

      The number of dimensions of the array.

      Returns number

    • get ref(): this

      Access an array by reference, incrementing the reference count.

      jax-js handles freeing arrays by using "move" semantics, like in Rust/C++. Whenever you pass an array into a function, that function should consume the array, and it will no longer be usable. For example, if you had:

      const x = np.array([1, 2, 3]);
      const y = np.add(x, x);

      The second line does not work because the first parameter consumes x, and then the second parameter will already have been freed / disposed.

      To fix this, you can write:

      const y = np.add(x.ref, x);
      

      Under the hood, every access to .ref increments the internal reference count of the array. The reference count starts at 1. When it hits 0, the memory behind the array is freed.

      Returns this

    • get shape(): number[]

      The shape of the array.

      Returns number[]

    • get size(): number

      The total number of elements in the array.

      Returns number

    • get weakType(): boolean

      Whether the array is weakly typed.

      Weakly typed arrays will cast to the dtype of the other operand. See promoteTypes() for details.

      Returns boolean

    Methods

    • Returns number

    • Iterate over the first dimension of this array, returning slices.

      This can be used to destructure arrays. For example:

      let x = np.array([[1, 2], [3, 4]]);
      let [a, b] = x;
      console.log(a.js()); // [1, 2]
      console.log(b.js()); // [3, 4]

      Returns IterableIterator<Array>

    • Convert this array into a primitive value.

      This only works for scalars (0-dimensional arrays). It lets you get values "out" of the JAX system. For instance, if x = np.array(5), then you can evaluate x + 1 and x ** 2 to get 6 and 25, respectively.

      This method is also called for == equality.

      Returns any

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Copy the array and cast to a specified dtype.

      Parameters

      Returns this

    • Wait for this array to finish evaluation.

      Operations and data loading in jax-js are lazy, so this function ensures that pending operations are dispatched and fully executed before it returns.

      If you are mapping from data() or dataSync(), it will also trigger dispatch of operations as well.

      Note: jax.blockUntilReady() is a higher-level API, it calls this asynchronously for multiple arrays.

      Returns Promise<Array>

    • Realize the array and return it as data.

      Returns Promise<
          | Float16Array<ArrayBuffer>
          | Float32Array<ArrayBuffer>
          | Int32Array<ArrayBuffer>
          | Uint32Array<ArrayBuffer>,
      >

    • Realize the array and return it as data. This is a sync variant and not recommended for performance reasons, as it will block rendering.

      Returns
          | Float16Array<ArrayBuffer>
          | Float32Array<ArrayBuffer>
          | Int32Array<ArrayBuffer>
          | Uint32Array<ArrayBuffer>

    • Return specified diagonals. See numpy.diagonal for full docs.

      Parameters

      • offset: number = 0
      • axis1: number = 0
      • axis2: number = 1

      Returns this

    • Manually decrement the reference count of the array.

      Arrays are created with reference count 1. Whenever it is used as argument to a function or other operation, it is disposed (i.e., reference count decreases by 1) automatically. Whenever a .ref is created, the reference count increases.

      You generally don't need to call this function directly since arrays are automatically disposed after being passed into an operation. One common exception is when writing a function and ignoring one of its arguments. In that case, by convention you should dispose of that argument manually.

      function myCustomOperation(a: np.Array, b: np.Array) {
      b.dispose(); // Needed to satisfy "move" rules.
      return a.add(1);
      }

      Returns void

    • Divide an array by this one.

      Parameters

      • other: Array | TracerValue

      Returns this

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Flatten the array without changing its data.

      Returns this

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Copy an element of an array to a numeric scalar and return it.

      Throws an error if the array does not have a single element. The array must either be rank-0, or all dimensions of the shape are 1.

      Returns number

    • Convert this array into a JavaScript object.

      This is a blocking operation that will compile all of the shaders and wait for execution to complete, synchronously. No other JavaScript code on the site will be run during shader execution.

      To avoid blocking, prefer jsAsync() when possible.

      Returns any

    • Convert this array into a JavaScript object, asynchronously.

      Returns Promise<any>

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Compute the average of the array elements along the specified axis.

      Parameters

      • axis: null | number | number[] = null
      • Optionalopts: { keepdims?: boolean }

      Returns Array

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Returns Array

    • Parameters

      • other: Array | TracerValue

      Returns Array

    • Product of the array elements over a given axis.

      Parameters

      • axis: null | number | number[] = null
      • Optionalopts: { keepdims?: boolean }

      Returns Array

    • Flatten the array without changing its data.

      Returns this

    • Give a new shape to an array without changing its data.

      One shape dimension can be -1. In this case, the value is inferred from the length of the array and remaining dimensions.

      Parameters

      • shape: number | number[]

      Returns this

    • Slice an array along one or more axes.

      This is the equivalent of slicing in Python, e.g. x[1:3, 2, :, None]. To mimic this in JavaScript, we would write:

      x.slice([1, 3], 2, [], null);
      

      The slice method accepts a variable number of arguments, each of which can be a number, an empty array, a single-element array, a two-element array, or null. The arguments are interpreted as follows:

      • A number n means to access the n-th element along that axis, removing that axis from the resulting shape.
      • An empty array [] means to keep that axis as-is, like : in Python.
      • A single-element array [i] means to start slicing from index i (inclusive) to the end of the axis, like x[i:].
      • A two-element array [i, j] means to slice from index i (inclusive) to index j (exclusive), like x[i:j].
      • null means to add a new axis at that position, like np.newaxis.

      Like in Python, negative indices are supported, which count from the end of the axis. For example, -1 means the last element.

      Strided slices are not yet implemented, so you cannot write x[::2] or similar.

      Advanced indexing by integer arrays is also supported. This translates to the "gather" primitive, and it allows you to access specific elements of the array by integer indices stored in another array.

      Parameters

      • ...index: (number | [] | [number, number] | Tracer | [number] | null)[]

      Returns this

    • Subtract an array from this one.

      Parameters

      • other: Array | TracerValue

      Returns this

    • Sum of the elements of the array over a given axis, or axes.

      Parameters

      • axis: null | number | number[] = null
      • Optionalopts: { keepdims?: boolean }

      Returns Array

    • Return a simple string representation of the array's dimensions.

      Returns string

    • Permute the dimensions of an array. Defaults to reversing the axis order.

      Parameters

      • Optionalperm: number[]

      Returns this