JaxLike

class probly.representation.jax_like.JaxLike(*args, **kwargs)

Bases: ArrayLike, Protocol, Generic

Protocol for array-like objects that implement JAX-specific APIs.

Every member declared here also exists on a real jax.Array, otherwise isinstance(jnp.ones(3), JaxLike) would silently be False. That rules out __jax_function__, which is declared on JaxLikeImplementation instead.

property T: Any

Transposed array.

astype(dtype: DTypeLike | None, copy: bool = False, device: jax.Device | Sharding | None = None) → Self

Return a copy of the array cast to the given data type.

property at: Any

Indexed update helper, see jax.Array.at.

block_until_ready() → Self

Block until the asynchronous computation of the array has finished.

copy() → Self

Return a copy of the array.

property device: Any

Device of the array.

devices() → set[Device]

Return the set of devices the array lives on.

property dtype: Any

Data type of the array.

flatten(order: str = 'C') → Self

Return a flattened version of the array.

item(*args: int) → bool | int | float | complex

Return the array as a Python scalar.

property mT: Any

Matrix transposed array.

property ndim: int

Number of dimensions.

ravel(order: str = 'C') → Self

Return a flattened version of the array.

reshape(*args: int | Sequence[int], order: str = 'C') → Self

Return a reshaped version of the array.

property shape: tuple[int, ...]

Shape of the array.

property sharding: Any

The sharding of the underlying array.

property size: int

Number of elements in the array.

squeeze(axis: int | Sequence[int] | None = None) → Self

Return the array with axes of length one removed.

to_device(device: Literal['cpu'], /, *, stream: int | Any | None = None) → Self

Move the array to a device.

tolist() → Any

Return the array as a (nested) Python list.

transpose(*args: int | Sequence[int] | None) → Self

Return a transposed version of the array.

view(dtype: DTypeLike | None = None, type: None = None) → Self

Return a bit-cast view of the array.