JaxDistanceBasedCredalSet

class probly.representation.credal_set.jax.JaxDistanceBasedCredalSet(nominal: JaxCategoricalDistribution, radius: Array)

Bases: JaxAxisProtected[Any], JaxCategoricalCredalSet, DistanceBasedCredalSet

Distance-based credal set around a nominal categorical distribution.

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

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 nominal distribution as the barycenter of the credal set.

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 a credal set from categorical distribution 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() → Array

Compute the lower envelope of the credal set.

For L1/TV distance, the tightest element-wise lower bound is max(0, nominal - radius).

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.

nominal: JaxCategoricalDistribution

property num_classes: int

Return the number of classes in the credal set.

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]] = {'nominal': 0, 'radius': 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.

radius: jax.Array

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.

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() → Array

Compute the upper envelope of the credal set.

For L1/TV distance, the tightest element-wise upper bound is min(1, nominal + radius).

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.