nami.core.specs

Functions

as_tuple(x)	normaliser to take flexible input and return always a tuple for convenience.
event_numel(event_shape)	returns the total number of elements in the event shape
flatten_event(x, event_ndim)	collapse all event dimensions into a single flat dimension
split_event(x, event_ndim)	given a tensor, return shape split into leading shape and event shape
unflatten_event(x, event_shape)	inverse of flatten_event
validate_shapes(tensor, event_ndim[, ...])	Runtime assertion helper to enforce explicit shapes and prevent silent broadcasting

Classes

TensorSpec(event_shape[, dtype])

Minimal tensor specification for models, samplers, and distributions.

class nami.core.specs.TensorSpec(event_shape, dtype=None)

Bases: object

Minimal tensor specification for models, samplers, and distributions.

Attributes:

event_shape (tuple[int, …]): The shape of a single event (sample, vector, matrix, etc). dtype (torch.dtype | None): The expected data type of the tensor.

dtype: dtype | None = None

property event_ndim: int

event_shape: tuple[int, ...]

property numel: int

Parameters:

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

• dtype (dtype | None)

nami.core.specs.as_tuple(x)

normaliser to take flexible input and return always a tuple for convenience.

Return type:

tuple[int, ...]

Parameters:

x (Iterable[int] | int | None)

nami.core.specs.event_numel(event_shape)

returns the total number of elements in the event shape

Return type:

int

Parameters:

event_shape (Iterable[int] | None)

nami.core.specs.flatten_event(x, event_ndim)

collapse all event dimensions into a single flat dimension

Return type:

Tensor

Parameters:

• x (Tensor)

• event_ndim (int)

nami.core.specs.split_event(x, event_ndim)

given a tensor, return shape split into leading shape and event shape

Return type:

tuple[tuple[int, ...], tuple[int, ...]]

Parameters:

• x (Tensor)

• event_ndim (int)

nami.core.specs.unflatten_event(x, event_shape)

inverse of flatten_event

Return type:

Tensor

Parameters:

• x (Tensor)

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

nami.core.specs.validate_shapes(tensor, event_ndim, expected_event_shape=None, batch_shape=None)

Runtime assertion helper to enforce explicit shapes and prevent silent broadcasting

Return type:

None

Parameters:

• tensor (Tensor)

• event_ndim (int)

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

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