JaxDirichletLevelSetCredalSet

class probly.representation.credal_set.jax.JaxDirichletLevelSetCredalSet(alphas: Array, threshold: Array)

Bases: JaxAxisProtected[Any], JaxCategoricalCredalSet, DirichletLevelSetCredalSet

Dirichlet density level set credal set.

Contains all distributions whose Dirichlet likelihood is at least threshold times the peak density.

property T: Self

Inverts the order of the axes of the underlying array.

alphas: 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.

property barycenter: JaxCategoricalDistribution

Return the Dirichlet mean as the barycenter.

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.

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

Return a flattened version of the array.

Parameters:

order – The read/write element order.

Returns:

The flattened array.

classmethod from_jax_sample(sample: JaxSample[JaxCategoricalDistribution]) → Self

Create from a jax sample (not supported).

Raises:

NotImplementedError – Always, as this credal set type cannot be created from samples.

classmethod from_sample(sample: Sample[JaxCategoricalDistribution]) → Self

Create a credal set from a finite sample.

Parameters:

sample – The sample to create the credal set from.

Returns:

The created credal set.

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.

lower(key: Array | None = None) → Array

Compute per-class lower bounds via Monte Carlo sampling.

Parameters:

key – PRNG key used for sampling. Defaults to a fixed key for determinism.

Returns:

Lower probability bounds, shape (…, K).

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

Return the number of classes.

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

permitted_functions = {}

classmethod primary_protected_name() → str

Return the first protected field (dict order).

protected_axes: ClassVar[dict[str, int]] = {'alphas': 1, 'threshold': 0}

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.

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.

threshold: jax.Array

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'

upper(key: Array | None = None) → Array

Compute per-class upper bounds via Monte Carlo sampling.

Parameters:

key – PRNG key used for sampling. Defaults to a fixed key for determinism.

Returns:

Upper probability bounds, shape (…, K).

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, JaxProtectedValue], func: Callable | None = None) → JaxAxisProtected[J]

Return a copy with updated protected field values.