JaxLogitCategoricalDistribution

class probly.representation.distribution.jax_categorical.JaxLogitCategoricalDistribution(array: Array)

Bases: JaxCategoricalDistribution

A categorical distribution represented by logits.

property T: Self

Inverts the order of the axes of the underlying array.

array: jax.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.

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.

property at: JaxAxisProtectedIndexUpdateHelper[J]

Return the out-of-place indexed update helper, mirroring jax.Array.at.

block_until_ready() → Self

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

Returns:

The object itself.

copy() → Self

Return a copy of the array.

property device: Device

Device of the array.

devices() → set[Device]

Return the set of devices the array lives on.

property dtype: dtype

Data type of the array.

entropy() → ArrayLike

Compute entropy.

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

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

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 log_probabilities: Array

Get the log probabilities of the categorical distribution.

property logits: Array

Get logits of the categorical distribution.

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.

property num_classes: int

Get the number of classes.

numpy(*, force: bool = False) → ndarray

permitted_functions: ClassVar[set[Callable]] = {<function jax_average>, <function jax_mean>}

classmethod primary_protected_name() → str

Return the first protected field (dict order).

property probabilities: Array

Get the probabilities of the categorical distribution.

protected_axes: ClassVar[dict[str, int]] = {'array': 1}

property protected_shape: tuple[int, ...]

Protected trailing shape of the primary field.

protected_value() → JaxProtectedValue

Return the primary protected value.

protected_values(func: Callable | None = None) → dict[str, JaxProtectedValue] | None

Return all protected field values as-is.

Optionally takes the jax function that triggered the call for context. This can be used to conditionally modify the returned values or prevent them from being accessed.

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

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

Return a copy with reshaped protected values.

sample(num_samples: int = 1, prng_key: ArrayLike | None = None) → JaxSample[jax.Array]

Sample from the categorical distribution (Jax backend).

Parameters:

• num_samples – The number of samples to draw.

• prng_key – The PRNG key to sample with. Defaults to a fresh key seeded from OS entropy, so repeated calls draw different samples.

Returns:

The drawn class indices with the sample axis first.

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.

Parameters:

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

Returns:

The squeezed array.

stop_gradient() → Self

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.

take_along_axis(indices: Array, axis: int = -1) → Self

Return a copy with gathered protected values along a batch dimension.

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

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

Return the array as a (nested) Python list.

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

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]

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

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:

• aux_data – The auxiliary data produced by tree_flatten().

• children – The children produced by tree_flatten().

Returns:

The reconstructed object.

type = 'categorical'

property unnormalized_probabilities: Array

Get unnormalized probabilities of the categorical distribution.

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

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.

with_protected_values(values: dict[str, Any], func: Callable | None = None) → JaxAxisProtected[jax.Array]

Return a copy with a replaced primary protected value.