skwdro.distributions package

Submodules

skwdro.distributions.dirac_distribution module

class skwdro.distributions.dirac_distribution.Dirac(loc: Tensor, n_batch_dims: int = 0, validate_args: bool | None = None)

Bases: ExponentialFamily

property arg_constraints: Dict[str, Constraint]

Returns a dictionary from argument names to Constraint objects that should be satisfied by each argument of this distribution. Args that are not tensors need not appear in this dict.

entropy() → Tensor

Method to compute the entropy using Bregman divergence of the log normalizer.

enumerate_support(expand: bool = True) → Tensor

Returns tensor containing all values supported by a discrete distribution. The result will enumerate over dimension 0, so the shape of the result will be (cardinality,) + batch_shape + event_shape (where event_shape = () for univariate distributions).

Note that this enumerates over all batched tensors in lock-step [[0, 0], [1, 1], …]. With expand=False, enumeration happens along dim 0, but with the remaining batch dimensions being singleton dimensions, [[0], [1], ...

To iterate over the full Cartesian product use itertools.product(m.enumerate_support()).

Args:

expand (bool): whether to expand the support over the

batch dims to match the distribution’s batch_shape.

Returns:

Tensor iterating over dimension 0.

expand(batch_shape: Size | List[int] | Tuple[int, ...], _instance=None)

Returns a new distribution instance (or populates an existing instance provided by a derived class) with batch dimensions expanded to batch_shape. This method calls expand on the distribution’s parameters. As such, this does not allocate new memory for the expanded distribution instance. Additionally, this does not repeat any args checking or parameter broadcasting in __init__.py, when an instance is first created.

Parameters:

batch_shape: torch.Size

the desired expanded size.

_instance: None|Dirac

new instance provided by subclasses that need to override .expand.

Returns:

new: Dirac

New distribution instance with batch dimensions expanded to batch_size.

has_rsample = True

log_prob(value: Tensor) → Tensor

Returns the log of the probability density/mass function evaluated at value.

Args:

value (Tensor):

property mean: Tensor

Returns the mean of the distribution.

property mode: Tensor

Returns the mode of the distribution.

perplexity() → Tensor

Returns perplexity of distribution, batched over batch_shape.

Returns:

Tensor of shape batch_shape.

rsample(sample_shape: Size | List[int] | Tuple[int, ...] = ()) → Tensor

Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

property support: Constraint | None

Returns a Constraint object representing this distribution’s support.

property variance: Tensor

Returns the variance of the distribution.

skwdro.distributions.dirac_distribution.cast_to_size(shape: Size | List[int] | Tuple[int, ...]) → Size

Module contents