probly.representation.jax_like.JaxLikeImplementation

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

Bases: ArrayLike[Any], ABC

Protocol for array-like objects that behave like JAX arrays.

property T: Any

Transposed array.

property device: Any

Device of the array.

property dtype: Any

Data type of the array.

property mT: Any

Matrix transposed array.

property ndim: int

Number of dimensions.

property shape: tuple[int, ...]

Shape of the array.

property size: int

Number of elements in the array.

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

Move the array to a device.