ShapeMixin

class ShapeMixin(shape: tuple[int, ...] | int)

Bases: BatchMixin

Mixin for modules with a concept of shape.

This mixin does not provide options for altering shape, only storing the given shape and offering related properties.

Parameters:

shape (tuple[int, ...] | int) – shape of the group being represented, excluding batch size.

property count: int

Number of elements in the module, excluding replication along the batch dim.

Returns:

number of elements in the module.

Return type:

int

property shape: tuple[int, ...]

Shape of the module.

Returns:

Shape of the module.

Return type:

tuple[int, …]