JaxIntervalConformalSet

class probly.representation.conformal_set.jax.JaxIntervalConformalSet(array: Array)

Bases: JaxAxisProtected[Any], IntervalConformalSet

Interval conformal set backed by a JAX array storing lower and upper bounds.

property T: Self

Inverts the order of the axes of the underlying array.

array: jnp.ndarray

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.

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_arrays(lower: Array, upper: Array) → Self

Create an interval conformal set from lower and upper bound arrays.

Parameters:

• lower – The lower bound array.

• upper – The upper bound array.

Returns:

The created interval conformal set.

classmethod from_samples(lower: JaxSample, upper: JaxSample) → Self

Create an interval conformal set from two JaxSamples.

Parameters:

• lower – The lower bound sample.

• upper – The upper bound sample.

Returns:

The created interval conformal 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.

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.

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]] = {'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.

property set_size: Array

Return the sizes of the conformal sets.

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 = 'interval'

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.