nami.distributions.normal

Classes

DiagonalNormal(loc, scale, *[, event_ndim, ...])
StandardNormal(event_shape, *[, ...])

class nami.distributions.normal.DiagonalNormal(loc, scale, *, event_ndim=1, validate_args=True)

Bases: Distribution

Parameters:

• loc (torch.Tensor)

• scale (torch.Tensor)

• event_ndim (int)

• validate_args (bool)

arg_constraints: ClassVar[dict[str, object]] = {}

expand(batch_shape, _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 (Size) – the desired expanded size.

• _instance (Distribution | None) – new instance provided by subclasses that need to override .expand.

Return type:

DiagonalNormal

Returns:

New distribution instance with batch dimensions expanded to batch_size.

has_rsample = True

log_prob(value)

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

Parameters:

value (Tensor)

Return type:

Tensor

property mean: Tensor

Returns the mean of the distribution.

rsample(sample_shape=())

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

Return type:

Tensor

Parameters:

sample_shape (Size)

sample(sample_shape=())

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

Return type:

Tensor

Parameters:

sample_shape (Size)

property variance: Tensor

Returns the variance of the distribution.

class nami.distributions.normal.StandardNormal(event_shape, *, batch_shape=None, device=None, dtype=None, validate_args=True)

Bases: Distribution

Parameters:

• event_shape (tuple[int, ...] | int)

• batch_shape (tuple[int, ...] | int | None)

• device (torch.device | None)

• dtype (torch.dtype | None)

• validate_args (bool)

arg_constraints: ClassVar[dict[str, object]] = {}

expand(batch_shape, _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 (Size) – the desired expanded size.

• _instance (Distribution | None) – new instance provided by subclasses that need to override .expand.

Return type:

StandardNormal

Returns:

New distribution instance with batch dimensions expanded to batch_size.

has_rsample = True

log_prob(value)

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

Parameters:

value (Tensor)

Return type:

Tensor

property mean: Tensor

Returns the mean of the distribution.

rsample(sample_shape=())

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

Return type:

Tensor

Parameters:

sample_shape (Size)

sample(sample_shape=())

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

Return type:

Tensor

Parameters:

sample_shape (Size)

property variance: Tensor

Returns the variance of the distribution.