JaxCategoricalDistributionSample

class probly.representation.distribution.jax_categorical.JaxCategoricalDistributionSample(array: D, sample_axis: int, weights: jax.Array | None = None)[source]

Bases: CategoricalDistributionSample[JaxCategoricalDistribution], JaxSample[JaxCategoricalDistribution]

Sample type for empirical second-order categorical distributions.

property T: Self

Inverts the order of the axes of the underlying array.

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

Cast the underlying array to a new 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 new JaxSample with the cast array.

property at: Any

The indexed update helper of the underlying array.

block_until_ready() → Self[source]

Block until the asynchronous computation of the underlying arrays has finished.

Returns:

The sample itself.

concat(other: Sample[D]) → Self[source]

Append another sample to this sample.

copy() → Self[source]

Create a copy of the JaxSample.

Returns:

A copy of the JaxSample.

property device: Any

The device of the underlying array.

property dtype: DTypeLike

The data type of the underlying 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.

classmethod from_iterable(samples: Iterable[D], weights: Iterable[float] | None = None, sample_axis: SampleAxis = 'auto', dtype: DTypeLike | None = None) → Self[source]

Create an JaxSample from a sequence of samples.

Parameters:

  • samples – The predictions to create the sample from.

  • weights – Optional weights for the samples.

  • sample_axis – The dimension along which samples are organized.

  • dtype – Desired data type of the array.

Returns:

The created JaxSample.

classmethod from_sample(sample: Sample[D], sample_axis: SampleAxis = 'auto', dtype: DTypeLike | None = None) → Self[source]

Create a new Sample from an existing Sample.

Parameters:

  • sample – The sample to create the new sample from.

  • sample_axis – The dimension along which samples are organized.

Returns:

The created Sample.

property is_weighted: bool

Return whether the samples are weighted.

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.

move_sample_axis(new_sample_axis: int) → JaxSample[D][source]

Return a new JaxSample with the sample dimension moved to new_sample_axis.

Parameters:

new_sample_axis – The new sample dimension.

Returns:

A new JaxSample with the sample dimension moved.

property ndim: int

The number of dimensions of the underlying array.

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.

sample_axis: int
sample_mean() → D[source]

Compute the mean of the sample.

property sample_size: int

Return the number of samples.

sample_space[source]

alias of JaxCategoricalDistribution

sample_std(ddof: int = 0) → D[source]

Compute the standard deviation of the sample.

sample_var(ddof: int = 0) → D[source]

Compute the variance of the sample.

property samples: D

Return an iterator over the samples.

property shape: tuple[int, ...]

The shape of the underlying array.

property size: int

The total number of elements in the underlying 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: Device | Sharding, /, *, stream: int | Any | None = None) → Self[source]

Move the underlying array to the specified device.

Parameters:

  • device – The target device.

  • stream – not implemented, passing a non-None value will lead to an error.

Returns:

A new JaxSample on the specified device.

Raises:

NotImplementedError – If a stream is given.

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.

weights: Array | None