meridian.backend.Tensor

Array base class for JAX

jax.Array is the public interface for instance checks and type annotation of JAX arrays and tracers. Its main applications are in instance checks and type annotations; for example::

x = jnp.arange(5) isinstance(x, jax.Array) # returns True both inside and outside traced functions.

def f(x: Array) -> Array: # type annotations are valid for traced and non-traced types. return x

jax.Array should not be used directly for creation of arrays; instead you should use array creation routines offered in :mod:jax.numpy, such as :func:jax.numpy.array, :func:jax.numpy.zeros, :func:jax.numpy.ones, :func:jax.numpy.full, :func:jax.numpy.arange, etc.

T Compute the all-axis array transpose.

Refer to :func:jax.numpy.transpose for details.

addressable_shards List of addressable shards.
at Helper property for index update functionality.

The at property provides a functionally pure equivalent of in-place array modifications.

============================== ================================ Alternate syntax Equivalent In-place expression ============================== ================================ x = x.at[idx].set(y) x[idx] = y x = x.at[idx].add(y) x[idx] += y x = x.at[idx].subtract(y) x[idx] -= y x = x.at[idx].multiply(y) x[idx] *= y x = x.at[idx].divide(y) x[idx] /= y x = x.at[idx].power(y) x[idx] **= y x = x.at[idx].min(y) x[idx] = minimum(x[idx], y) x = x.at[idx].max(y) x[idx] = maximum(x[idx], y) x = x.at[idx].apply(ufunc) ufunc.at(x, idx) x = x.at[idx].get() x = x[idx] ============================== ================================

None of the x.at expressions modify the original x; instead they return a modified copy of x. However, inside a :py:func:~jax.jit compiled function, expressions like :code:x = x.at[idx].set(y) are guaranteed to be applied in-place.

Unlike NumPy in-place operations such as :code:x[idx] += y, if multiple indices refer to the same location, all updates will be applied (NumPy would only apply the last update, rather than applying all updates.) The order in which conflicting updates are applied is implementation-defined and may be nondeterministic (e.g., due to concurrency on some hardware platforms).

By default, JAX assumes that all indices are in-bounds. Alternative out-of-bound index semantics can be specified via the mode parameter (see below).

committed Whether the array is committed or not.

An array is committed when it is explicitly placed on device(s) via JAX APIs. For example, jax.device_put(np.arange(8), jax.devices()[0]) is committed to device 0. While jax.device_put(np.arange(8)) is uncommitted and will be placed on the default device.

Computations involving some committed inputs will happen on the committed device(s) and the result will be committed on the same device(s). Invoking an operation on arguments that are committed to different device(s) will raise an error.

device Array API-compatible device attribute.

For single-device arrays, this returns a Device. For sharded arrays, this returns a Sharding.

dtype The data type (:class:numpy.dtype) of the array.
flat Not implemented: Use :meth:~jax.Array.flatten instead.
global_shards List of global shards.
imag Return the imaginary part of the array.
is_fully_addressable Is this Array fully addressable?

A jax.Array is fully addressable if the current process can address all of the devices named in the :class:Sharding. is_fully_addressable is equivalent to "is_local" in multi-process JAX.

Note that fully replicated is not equal to fully addressable i.e. a jax.Array which is fully replicated can span across multiple hosts and is not fully addressable.

is_fully_replicated Is this Array fully replicated?
itemsize Length of one array element in bytes.
mT Compute the (batched) matrix transpose.

Refer to :func:jax.numpy.matrix_transpose for details.

nbytes Total bytes consumed by the elements of the array.
ndim The number of dimensions in the array.
real Return the real part of the array.
shape The shape of the array.
sharding The sharding for the array.
size The total number of elements in the array.

Methods

addressable_data

addressable_data(
    index: int
) -> Array

Return an array of the addressable data at a particular index.

all

all(
    axis: reductions.Axis = None,
    out: None = None,
    keepdims: bool = False,
    *,
    where: (ArrayLike | None) = None
) -> Array

Test whether all array elements along a given axis evaluate to True.

Refer to :func:jax.numpy.all for the full documentation.

any

any(
    axis: reductions.Axis = None,
    out: None = None,
    keepdims: bool = False,
    *,
    where: (ArrayLike | None) = None
) -> Array

Test whether any array elements along a given axis evaluate to True.

Refer to :func:jax.numpy.any for the full documentation.

argmax

argmax(
    axis: (int | None) = None, out: None = None, keepdims: (bool | None) = None
) -> Array

Return the index of the maximum value.

Refer to :func:jax.numpy.argmax for the full documentation.

argmin

argmin(
    axis: (int | None) = None, out: None = None, keepdims: (bool | None) = None
) -> Array

Return the index of the minimum value.

Refer to :func:jax.numpy.argmin for the full documentation.

argpartition

argpartition(
    kth: int, axis: int = -1
) -> Array

Return the indices that partially sort the array.

Refer to :func:jax.numpy.argpartition for the full documentation.

argsort

argsort(
    axis: (int | None) = -1,
    *,
    kind: None = None,
    order: None = None,
    stable: bool = True,
    descending: bool = False
) -> Array

Return the indices that sort the array.

Refer to :func:jax.numpy.argsort for the full documentation.

astype

astype(
    dtype: (DTypeLike | None),
    copy: bool = False,
    device: (xc.Device | Sharding | None) = None
) -> Array

Copy the array and cast to a specified dtype.

This is implemented via :func:jax.lax.convert_element_type, which may have slightly different behavior than :meth:numpy.ndarray.astype in some cases. In particular, the details of float-to-int and int-to-float casts are implementation dependent.

byteswap

byteswap() -> Array

Swap the bytes of the array elements.

This switches between a little-endian and big-endian data representation.

Returns
An array with the same dtype as self, with underlying bytes of each entry reversed.

Examples
```

import jax.numpy as jnp x = jnp.arange(5, dtype='int32') x Array([0, 1, 2, 3, 4], dtype=int32) x.byteswap() Array([ 0, 16777216, 33554432, 50331648, 67108864], dtype=int32) ```

When the resulting bytes are viewed as a big-endian dtype (possible in NumPy, but not in JAX) they represent the original values:

>>> import numpy as np
>>> np.array(x.byteswap()).view('>i4')  # view as big-endian
array([0, 1, 2, 3, 4], dtype='>i4')

Calling byteswap twice will return the original array:

>>> x.byteswap().byteswap()
Array([0, 1, 2, 3, 4], dtype=int32)

choose

choose(
    choices: Sequence[ArrayLike],
    out: None = None,
    mode: str = 'raise'
) -> Array

Construct an array choosing from elements of multiple arrays.

Refer to :func:jax.numpy.choose for the full documentation.

clip

clip(
    min: (ArrayLike | None) = None, max: (ArrayLike | None) = None
) -> Array

Return an array whose values are limited to a specified range.

Refer to :func:jax.numpy.clip for full documentation.

compress

compress(
    condition: ArrayLike,
    axis: (int | None) = None,
    *,
    out: None = None,
    size: (int | None) = None,
    fill_value: ArrayLike = 0
) -> Array

Return selected slices of this array along given axis.

Refer to :func:jax.numpy.compress for full documentation.

conj

conj() -> Array

Return the complex conjugate of the array.

Refer to :func:jax.numpy.conj for the full documentation.

conjugate

conjugate() -> Array

Return the complex conjugate of the array.

Refer to :func:jax.numpy.conjugate for the full documentation.

copy

copy() -> Array

Return a copy of the array.

Refer to :func:jax.numpy.copy for the full documentation.

copy_to_host_async

copy_to_host_async()

Copies an Array to the host asynchronously.

For arrays that live an an accelerator, such as a GPU or a TPU, JAX may cache the value of the array on the host. Normally this happens behind the scenes when the value of an on-device array is requested by the user, but waiting to initiate a device-to-host copy until the value is requested requires that JAX block the caller while waiting for the copy to complete.

copy_to_host_async requests that JAX populate its on-host cache of an array, but does not wait for the copy to complete. This may speed up a future on-host access to the array's contents.

cumprod

cumprod(
    axis: (int | None) = None,
    dtype: (DTypeLike | None) = None,
    out: None = None
) -> Array

Return the cumulative product of the array.

Refer to :func:jax.numpy.cumprod for the full documentation.

cumsum

cumsum(
    axis: (int | None) = None,
    dtype: (DTypeLike | None) = None,
    out: None = None
) -> Array

Return the cumulative sum of the array.

Refer to :func:jax.numpy.cumsum for the full documentation.

diagonal

diagonal(
    offset: int = 0, axis1: int = 0, axis2: int = 1
) -> Array

Return the specified diagonal from the array.

Refer to :func:jax.numpy.diagonal for the full documentation.

dot

dot(
    b: ArrayLike,
    *,
    precision: lax.PrecisionLike = None,
    preferred_element_type: (DTypeLike | None) = None
) -> Array

Compute the dot product of two arrays.

Refer to :func:jax.numpy.dot for the full documentation.

flatten

flatten(
    order: str = 'C', *, out_sharding=None
) -> Array

Flatten array into a 1-dimensional shape.

Refer to :func:jax.numpy.ravel for the full documentation.

item

item(
    *args
) -> (bool | int | float | complex)

Copy an element of an array to a standard Python scalar and return it.

max

max(
    axis: reductions.Axis = None,
    out: None = None,
    keepdims: bool = False,
    initial: (ArrayLike | None) = None,
    where: (ArrayLike | None) = None
) -> Array

Return the maximum of array elements along a given axis.

Refer to :func:jax.numpy.max for the full documentation.

mean

mean(
    axis: reductions.Axis = None,
    dtype: (DTypeLike | None) = None,
    out: None = None,
    keepdims: bool = False,
    *,
    where: (ArrayLike | None) = None
) -> Array

Return the mean of array elements along a given axis.

Refer to :func:jax.numpy.mean for the full documentation.

min

min(
    axis: reductions.Axis = None,
    out: None = None,
    keepdims: bool = False,
    initial: (ArrayLike | None) = None,
    where: (ArrayLike | None) = None
) -> Array

Return the minimum of array elements along a given axis.

Refer to :func:jax.numpy.min for the full documentation.

nonzero

nonzero(
    *,
    fill_value: (None | ArrayLike | tuple[ArrayLike, ...]) = None,
    size: (int | None) = None
) -> tuple[Array, ...]

Return indices of nonzero elements of an array.

Refer to :func:jax.numpy.nonzero for the full documentation.

prod

prod(
    axis: reductions.Axis = None,
    dtype: (DTypeLike | None) = None,
    out: None = None,
    keepdims: bool = False,
    initial: (ArrayLike | None) = None,
    where: (ArrayLike | None) = None,
    promote_integers: bool = True
) -> Array

Return product of the array elements over a given axis.

Refer to :func:jax.numpy.prod for the full documentation.

ptp

ptp(
    axis: reductions.Axis = None, out: None = None, keepdims: bool = False
) -> Array

Return the peak-to-peak range along a given axis.

Refer to :func:jax.numpy.ptp for the full documentation.

ravel

ravel(
    order: str = 'C', *, out_sharding=None
) -> Array

Flatten array into a 1-dimensional shape.

Refer to :func:jax.numpy.ravel for the full documentation.

repeat

repeat(
    repeats: ArrayLike,
    axis: (int | None) = None,
    *,
    total_repeat_length: (int | None) = None,
    out_sharding: (NamedSharding | PartitionSpec | None) = None
) -> Array

Construct an array from repeated elements.

Refer to :func:jax.numpy.repeat for the full documentation.

reshape

reshape(
    *args, order: str = 'C', out_sharding=None
) -> Array

Returns an array containing the same data with a new shape.

Refer to :func:jax.numpy.reshape for full documentation.

round

round(
    decimals: int = 0, out: None = None
) -> Array

Round array elements to a given decimal.

Refer to :func:jax.numpy.round for full documentation.

searchsorted

searchsorted(
    v: ArrayLike,
    side: str = 'left',
    sorter: (ArrayLike | None) = None,
    *,
    method: str = 'scan'
) -> Array

Perform a binary search within a sorted array.

Refer to :func:jax.numpy.searchsorted for full documentation.

sort

sort(
    axis: (int | None) = -1,
    *,
    kind: None = None,
    order: None = None,
    stable: bool = True,
    descending: bool = False
) -> Array

Return a sorted copy of an array.

Refer to :func:jax.numpy.sort for full documentation.

squeeze

squeeze(
    axis: reductions.Axis = None
) -> Array

Remove one or more length-1 axes from array.

Refer to :func:jax.numpy.squeeze for full documentation.

std

std(
    axis: reductions.Axis = None,
    dtype: (DTypeLike | None) = None,
    out: None = None,
    ddof: int = 0,
    keepdims: bool = False,
    *,
    where: (ArrayLike | None) = None,
    correction: (int | float | None) = None
) -> Array

Compute the standard deviation along a given axis.

Refer to :func:jax.numpy.std for full documentation.

sum

sum(
    axis: reductions.Axis = None,
    dtype: (DTypeLike | None) = None,
    out: None = None,
    keepdims: bool = False,
    initial: (ArrayLike | None) = None,
    where: (ArrayLike | None) = None,
    promote_integers: bool = True
) -> Array

Sum of the elements of the array over a given axis.

Refer to :func:jax.numpy.sum for full documentation.

swapaxes

swapaxes(
    axis1: int, axis2: int
) -> Array

Swap two axes of an array.

Refer to :func:jax.numpy.swapaxes for full documentation.

take

take(
    indices: ArrayLike,
    axis: (int | None) = None,
    out: None = None,
    mode: (str | None) = None,
    unique_indices: bool = False,
    indices_are_sorted: bool = False,
    fill_value: (StaticScalar | None) = None
) -> Array

Take elements from an array.

Refer to :func:jax.numpy.take for full documentation.

to_device

to_device(
    device: (xc.Device | Sharding), *, stream: (int | Any | None) = None
)

Return a copy of the array on the specified device

Args
device :class:~jax.Device or :class:~jax.sharding.Sharding to which the created array will be committed.
stream not implemented, passing a non-None value will lead to an error.

Returns
copy of array placed on the specified device or devices.

trace

trace(
    offset: (int | ArrayLike) = 0,
    axis1: int = 0,
    axis2: int = 1,
    dtype: (DTypeLike | None) = None,
    out: None = None
) -> Array

Return the sum along the diagonal.

Refer to :func:jax.numpy.trace for full documentation.

transpose

transpose(
    *args
) -> Array

Returns a copy of the array with axes transposed.

Refer to :func:jax.numpy.transpose for full documentation.

var

var(
    axis: reductions.Axis = None,
    dtype: (DTypeLike | None) = None,
    out: None = None,
    ddof: int = 0,
    keepdims: bool = False,
    *,
    where: (ArrayLike | None) = None,
    correction: (int | float | None) = None
) -> Array

Compute the variance along a given axis.

Refer to :func:jax.numpy.var for full documentation.

view

view(
    dtype: (DTypeLike | None) = None, type: None = None
) -> Array

Return a bitwise copy of the array, viewed as a new dtype.

This is fuller-featured wrapper around :func:jax.lax.bitcast_convert_type.

If the source and target dtype have the same bitwidth, the result has the same shape as the input array. If the bitwidth of the target dtype is different from the source, the size of the last axis of the result is adjusted accordingly.

>>> jnp.zeros([1,2,3], dtype=jnp.int16).view(jnp.int8).shape
(1, 2, 6)
>>> jnp.zeros([1,2,4], dtype=jnp.int8).view(jnp.int16).shape
(1, 2, 2)

Conversions involving booleans are not well-defined in all situations. With regards to the shape of result as explained above, booleans are treated as having a bitwidth of 8. However, when converting to a boolean array, the input should only contain 0 or 1 bytes. Otherwise, results may be unpredictable or may change depending on how the result is used.

This conversion is guaranteed and safe::

  >>> jnp.array([1, 0, 1], dtype=jnp.int8).view(jnp.bool_)
  Array([ True, False,  True], dtype=bool)

However, there are no guarantees about the results of any expression involving a view such as this: jnp.array([1, 2, 3], dtype=jnp.int8).view(jnp.bool_). In particular, the results may change between JAX releases and depending on the platform. To safely convert such an array to a boolean array, compare it with 0::

  >>> jnp.array([1, 2, 0], dtype=jnp.int8) != 0
  Array([ True,  True, False], dtype=bool)

Args
dtype An optional output dtype. If not specified, the output dtype is the same as the input dtype.
type Not implemented; accepted for NumPy compatibility.

Returns
The array, viewed as the new dtype. Unlike NumPy, the array may or may not be a copy of the input array.

__abs__

__abs__() -> Array

Alias of :func:jax.numpy.absolute.

__add__

__add__(
    other
)

__and__

__and__(
    other
)

__contains__

__contains__(
    other: ArrayLike
) -> Array

Implements contains for JAX arrays.

This is used by the Python in operator.

__eq__

__eq__(
    other
)

__floordiv__

__floordiv__(
    other
)

__ge__

__ge__(
    other
)

__getitem__

__getitem__(
    item
)

__gt__

__gt__(
    other
)

__invert__

__invert__() -> Array

Compute the bitwise inversion of an input.

JAX implementation of :func:numpy.invert. This function provides the implementation of the ~ operator for JAX arrays.

Args
x input array, must be boolean or integer typed.

Returns
An array of the same shape and dtype as `x, with the bits inverted.

See also
- :func:jax.numpy.bitwise_invert: Array API alias of this function. - :func:jax.numpy.logical_not: Invert after casting input to boolean.

Examples
```

x = jnp.arange(5, dtype='uint8') print(x) [0 1 2 3 4] print(jnp.invert(x)) [255 254 253 252 251] ```

This function implements the unary ~ operator for JAX arrays:

>>> print(~x)
[255 254 253 252 251]

:func:invert operates bitwise on the input, and so the meaning of its output may be more clear by showing the bitwise representation:

>>> with jnp.printoptions(formatter={'int': lambda x: format(x, '#010b')}):
...   print(f"{x  = }")
...   print(f"{~x = }")
x  = Array([0b00000000, 0b00000001, 0b00000010, 0b00000011, 0b00000100], dtype=uint8)
~x = Array([0b11111111, 0b11111110, 0b11111101, 0b11111100, 0b11111011], dtype=uint8)

For boolean inputs, :func:invert is equivalent to :func:logical_not:

>>> x = jnp.array([True, False, True, True, False])
>>> jnp.invert(x)
Array([False,  True, False, False,  True], dtype=bool)

__le__

__le__(
    other
)

__lshift__

__lshift__(
    other
)

__lt__

__lt__(
    other
)

__matmul__

__matmul__(
    other
)

__mod__

__mod__(
    other
)

__mul__

__mul__(
    other
)

__ne__

__ne__(
    other
)

__neg__

__neg__() -> Array

Return element-wise negative values of the input.

JAX implementation of :obj:numpy.negative.

Args
x input array or scalar.

Returns
An array with same shape and dtype as x containing -x.

See also
- :func:jax.numpy.positive: Returns element-wise positive values of the input. - :func:jax.numpy.sign: Returns element-wise indication of sign of the input.

Note
jnp.negative, when applied over unsigned integer, produces the result of their two's complement negation, which typically results in unexpected large positive values due to integer underflow.

Examples
For real-valued inputs:

>>> x = jnp.array([0., -3., 7])
>>> jnp.negative(x)
Array([-0.,  3., -7.], dtype=float32)

For complex inputs:

>>> x1 = jnp.array([1-2j, -3+4j, 5-6j])
>>> jnp.negative(x1)
Array([-1.+2.j,  3.-4.j, -5.+6.j], dtype=complex64)

For unit32:

>>> x2 = jnp.array([5, 0, -7]).astype(jnp.uint32)
>>> x2
Array([         5,          0, 4294967289], dtype=uint32)
>>> jnp.negative(x2)
Array([4294967291,          0,          7], dtype=uint32)

__or__

__or__(
    other
)

__pos__

__pos__() -> Array

Return element-wise positive values of the input.

JAX implementation of :obj:numpy.positive.

Args
x input array or scalar

Returns
An array of same shape and dtype as x containing +x.

Note
jnp.positive is equivalent to x.copy() and is defined only for the types that support arithmetic operations.

See also
- :func:jax.numpy.negative: Returns element-wise negative values of the input. - :func:jax.numpy.sign: Returns element-wise indication of sign of the input.

Examples
For real-valued inputs:

>>> x = jnp.array([-5, 4, 7., -9.5])
>>> jnp.positive(x)
Array([-5. ,  4. ,  7. , -9.5], dtype=float32)
>>> x.copy()
Array([-5. ,  4. ,  7. , -9.5], dtype=float32)

For complex inputs:

>>> x1 = jnp.array([1-2j, -3+4j, 5-6j])
>>> jnp.positive(x1)
Array([ 1.-2.j, -3.+4.j,  5.-6.j], dtype=complex64)
>>> x1.copy()
Array([ 1.-2.j, -3.+4.j,  5.-6.j], dtype=complex64)

For uint32:

>>> x2 = jnp.array([6, 0, -4]).astype(jnp.uint32)
>>> x2
Array([         6,          0, 4294967292], dtype=uint32)
>>> jnp.positive(x2)
Array([         6,          0, 4294967292], dtype=uint32)

__pow__

__pow__(
    other
)

__radd__

__radd__(
    other
)

__rand__

__rand__(
    other
)

__rfloordiv__

__rfloordiv__(
    other
)

__rlshift__

__rlshift__(
    other
)

__rmatmul__

__rmatmul__(
    other
)

__rmod__

__rmod__(
    other
)

__rmul__

__rmul__(
    other
)

__ror__

__ror__(
    other
)

__rpow__

__rpow__(
    other
)

__rrshift__

__rrshift__(
    other
)

__rshift__

__rshift__(
    other
)

__rsub__

__rsub__(
    other
)

__rtruediv__

__rtruediv__(
    other
)

__rxor__

__rxor__(
    other
)

__sub__

__sub__(
    other
)

__truediv__

__truediv__(
    other
)

__xor__

__xor__(
    other
)