TorchSparseLogCategoricalDistribution

class probly.representation.distribution.torch_sparse_log_categorical.TorchSparseLogCategoricalDistribution(group_ids: Tensor, entry_logits: Tensor)

Bases: TorchAxisProtected[Any], CategoricalDistribution[Tensor]

Sparse categorical distribution with grouped log-weights.

Shape: batch_shape. The final axis of group_ids and entry_logits stores sparse support entries. Entries with the same group id are combined by summing their exponentiated logits when converting to a dense categorical distribution.

property T: Self

Inverts the order of the dimensions of the underlying array.

clone(*, memory_format: memory_format = torch.preserve_format) → Self

Return a copy of the array.

cpu(memory_format: memory_format = torch.preserve_format) → Self

Move the array to the CPU.

cuda(device: device | str | None = None, non_blocking: bool = False, memory_format: memory_format = torch.preserve_format) → Self

Move the array to the GPU.

detach() → Self

Return a detached version of the array.

property device: device

Device of the array.

property dtype: dtype

Data type of the array.

entropy() → ArrayLike

Compute entropy.

entry_logits: torch.Tensor

gather(dim: int, index: Tensor) → Self

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

group_ids: torch.Tensor

property log_probabilities: Tensor

Get the log probabilities of the categorical distribution.

property logits: Tensor

Get logits of the categorical distribution.

property mH: Self

The adjoint (conjugate) transposed version of the underlying array.

property mT: Self

The transposed version of the underlying array.

property ndim: int

Number of dimensions.

property num_classes: int

Get the number of classes.

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

Convert dense probabilities to a numpy array.

permitted_functions = {}

permute(*dims: Size | int | tuple[int] | list[int]) → Self

Return a permuted version of the array.

classmethod primary_protected_name() → str

Return the first protected field (dict order).

property probabilities: Tensor

Get the probabilities of the categorical distribution.

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

property protected_shape: tuple[int, ...]

Protected trailing shape of the primary field.

protected_value() → TorchProtectedValue

Return the primary protected value.

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

Return all protected field values as-is.

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

reshape(*shape: int | tuple[int, ...]) → Self

Return a copy with reshaped protected values.

resolve_conj() → Self

Return a version of the array with any conjugate operations resolved.

resolve_neg() → Self

Return a version of the array with any negation operations resolved.

sample(num_samples: int = 1, rng: Generator | None = None) → TorchSample[Tensor]

Sample from the equivalent dense categorical distribution.

property shape: tuple[int, ...]

Shape of the array.

size(dim: int | None = None) → int | Size

Return the size of the array along the given dimension.

to(*args: Any, **kwargs: Any) → Self

Move and/or cast the tensor, mirroring torch.Tensor.to.

to_dense(num_classes: int | None = None) → TorchProbabilityCategoricalDistribution

Convert sparse grouped logits to a dense categorical distribution.

Parameters:

num_classes – Optional dense class count. Defaults to max(group_ids) + 1.

Returns:

Dense categorical distribution with zero mass for absent groups.

to_device(device: Literal['cpu'], /, *, stream: int | Any | None = None) → Self

Move the array to a device.

transpose(dim0: int, dim1: int) → Self

Return a transposed version of the array.

type = 'categorical'

uniform_logits() → TorchSparseLogCategoricalDistribution

Return a copy with identical groups and uniform sparse logits.

Returns:

A sparse categorical distribution with shared group_ids and zero logits.

property unnormalized_probabilities: Tensor

Get unnormalized probabilities of the categorical distribution.

with_protected_values(values: dict[str, TorchProtectedValue], func: Callable | None = None) → TorchAxisProtected[T]

Return a copy with updated protected field values.