quantax.utils.LogArray#

class quantax.utils.LogArray#

Log-amplitude representation of JAX arrays: value = sign * exp(logabs) where sign is \(\pm 1\) or a complex phase and logabs is real. Zero is encoded by sign=0, logabs=-inf.

The array is a PyTree with two leaves: sign and logabs. To convert it to a dense JAX array, use arr.value() or jnp.asarray(arr).

Warning

JAX doesn’t have a full support for customized arrays, so one should be careful when using LogArray. Here we list several possible problems.

1. Manipulations like jnp.fn(array) transform customized arrays to Array. To avoid it, call array.fn() whenever possible.

2. Computations like jax_array * customized_array always call jax_array.__mul__(customized_array), which returns a Array. To avoid it, use customized_array * jax_array.

__array__(dtype=None) → NDArray#: Convert to a numpy array.

__jax_array__() → Array#: Convert to a JAX array.

__neg__() → LogArray#: Negation of the represented value.

__abs__() → LogArray#: Absolute value of the represented array.

__mul__(other: ArrayLike) → LogArray#: Element-wise multiplication.

__rmul__(other: ArrayLike) → LogArray#: Reversed element-wise multiplication.

__truediv__(other: ArrayLike) → LogArray#: Element-wise division.

__rtruediv__(other: ArrayLike) → LogArray#: Reversed element-wise division.

__pow__(p: float | Array) → LogArray#: Element-wise power.

__add__(other: ArrayLike) → LogArray#: Element-wise addition.

__radd__(other: ArrayLike) → LogArray#: Reversed element-wise addition.

__sub__(other: ArrayLike) → LogArray#: Element-wise subtraction.

__rsub__(other: ArrayLike) → LogArray#: Reversed element-wise subtraction.

choose(*args, **kwargs) → LogArray#: Apply choose to sign and logabs component-wise.

compress(*args, **kwargs) → LogArray#: Apply compress to sign and logabs component-wise.

copy(*args, **kwargs) → LogArray#: Apply copy to sign and logabs component-wise.

diagonal(*args, **kwargs) → LogArray#: Apply diagonal to sign and logabs component-wise.

flatten(*args, **kwargs) → LogArray#: Apply flatten to sign and logabs component-wise.

ravel(*args, **kwargs) → LogArray#: Apply ravel to sign and logabs component-wise.

repeat(*args, **kwargs) → LogArray#: Apply repeat to sign and logabs component-wise.

reshape(*args, **kwargs) → LogArray#: Apply reshape to sign and logabs component-wise.

squeeze(*args, **kwargs) → LogArray#: Apply squeeze to sign and logabs component-wise.

swapaxes(*args, **kwargs) → LogArray#: Apply swapaxes to sign and logabs component-wise.

take(*args, **kwargs) → LogArray#: Apply take to sign and logabs component-wise.

transpose(*args, **kwargs) → LogArray#: Apply transpose to sign and logabs component-wise.

static from_value(x: ArrayLike) → LogArray#: Create from a JAX array / Python scalar.

property shape: tuple[int, ...]#: The shape of the represented array.

property dtype: str | type[Any] | dtype | SupportsDType#: The data type of the represented array.

property ndim: int#: The number of dimensions of the represented array.

property size: int#: The total number of elements in the represented array.

property nbytes: int#: The total number of bytes consumed by the represented array.

property sharding: Sharding#: The sharding of the represented array.

value() → Array#: Materialize the dense array value.

property T: LogArray#: Transpose of the represented array.

property mT: LogArray#: Matrix transpose of the represented array.

conj() → LogArray#: Complex conjugate of the represented array.

abs() → LogArray#: Absolute value of the represented array.

property real: LogArray#: Real part of the represented array.

property imag: LogArray#: Imaginary part of the represented array.

astype(dtype) → LogArray#: Cast the represented array to given dtype.

sum(axis: int | tuple[int, ...] | None = None, keepdims: bool = False) → LogArray#: Sum of array elements over a given axis.

mean(axis: int | tuple[int, ...] | None = None, keepdims: bool = False) → LogArray#: Mean of array elements over a given axis.

prod(axis: int | tuple[int, ...] | None = None, keepdims: bool = False) → LogArray#: Product of array elements over a given axis.