JaxLikeImplementation

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

Bases: ArrayLike, ABC, Generic

ABC implementation for array-like objects that behave like JAX arrays.

Subclasses only have to implement __jax_function__ plus the few members that cannot be expressed generically. Every concrete method below routes through the wrappers in probly.representation.jax_functions and thus back into __jax_function__.

Concrete (non-abstract) subclasses are additionally registered as JAX pytree nodes, which makes jax.jit, jax.vmap, jax.grad and jax.tree work on them. This does not replace the override machinery: jnp.sum(obj) still raises for a registered pytree.

at, the JAX stand-in for __setitem__, is left to subclasses: an indexed update has no single meaning across fields with different batch semantics.

property T: Self

Inverts the order of the axes of the underlying array.

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

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

Parameters:

  • dtype – The target data type.

  • copy – Whether to always return a copy.

  • device – The device the result should live on.

Returns:

A copy with every array-valued field cast to the given data type.

block_until_ready() → Self[source]

Block until the asynchronous computation of every array-valued field has finished.

Returns:

The object itself.

copy() → Self[source]

Return a copy of the array.

property device: Any

Device of the array.

property dtype: Any

Data type of the array.

flatten(order: str = 'C') → Self[source]

Return a flattened version of the array.

Parameters:

order – The read/write element order.

Returns:

The flattened array.

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

Return the array as a Python scalar.

Parameters:

*args – Optional index of the element to return.

Returns:

The selected element as a Python scalar.

property mH: Self

The adjoint (conjugate) transposed version of the underlying array.

property mT: Self

The version of the underlying array with its last two axes swapped.

property ndim: int

Number of dimensions.

ravel(order: str = 'C') → Self[source]

Return a flattened version of the array.

Parameters:

order – The read/write element order.

Returns:

The flattened array.

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

Return a reshaped version of the array.

Parameters:

  • *args – The target shape, either as a single sequence or as separate integers.

  • order – The read/write element order.

Returns:

The reshaped array.

property shape: tuple[int, ...]

Shape of the array.

property size: int

Number of elements in the array.

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

Return the array with axes of length one removed.

Parameters:

axis – The axis or axes to remove. If omitted, all length-one axes are removed.

Returns:

The squeezed array.

stop_gradient() → Self[source]

Return a copy detached from the autodiff graph.

Subclasses holding NumPy sidecar fields should override this to leave those fields untouched.

Returns:

A copy through which gradients do not propagate.

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

Move the array to the given device.

Parameters:

  • device – The target device.

  • stream – Unsupported by JAX, must be None.

Returns:

A copy with every array-valued field on the given device.

Raises:

NotImplementedError – If stream is not None.

tolist() → Any[source]

Return the array as a (nested) Python list.

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

Return a transposed version of the array.

Parameters:

*args – The permutation of the axes, either as a single sequence or as separate integers. If omitted, the axes are reversed.

Returns:

The transposed array.

tree_flatten() → tuple[tuple[Any, ...], Any][source]

Split the object into pytree children and static auxiliary data.

Fields holding arrays (anything exposing shape, plus None) become children, all remaining fields become auxiliary data. Subclasses may override this together with tree_unflatten().

Returns:

The children and the auxiliary data.

classmethod tree_unflatten(aux_data: Any, children: tuple[Any, ...]) → Self[source]

Rebuild an object from pytree children and auxiliary data.

The object is built without calling __init__, because JAX passes tracer objects here during tracing and those would fail the usual validation.

Parameters:
Returns:

The reconstructed object.

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

Return a bit-cast view of the array.

Parameters:

  • dtype – The data type to reinterpret the underlying bytes as.

  • type – Unsupported by JAX, must be None.

Returns:

A copy with every array-valued field bit-cast to the given data type.

Raises:

NotImplementedError – If type is not None.