bijx.ShapeInfo

class bijx.ShapeInfo[source]

Bases: object

Comprehensive shape information manager for array operations.

Manages the relationship between different types of array dimensions: event, spatial, channel, and batch dimensions. Automatically infers missing information from provided specifications.

The shape hierarchy follows: event_shape = space_shape + channel_shape This enables structured operations on multi-dimensional data like images, lattice fields, or other spatially-structured data with channels.

Parameters:

  • event_shape (tuple[int, ...] | None) – Complete event dimensions (spatial + channel).

  • space_shape (tuple[int, ...] | None) – Spatial dimensions only.

  • channel_shape (tuple[int, ...] | None) – Channel dimensions only.

  • event_dim (int | None) – Number of event dimensions.

  • space_dim (int | None) – Number of spatial dimensions.

  • channel_dim (int | None) – Number of channel dimensions.

Note

Only partial information needs to be provided. Missing information is automatically inferred from the provided specifications.

Example

>>> # For 2D images with RGB channels
>>> info = ShapeInfo(event_shape=(32, 32, 3), space_dim=2)
>>> info.space_shape  # Can be set via space_dim=2
(32, 32)
>>> info.channel_shape  # Can be set via channel_dim=1
(3,)
__init__(event_shape=None, space_shape=None, channel_shape=None, event_dim=None, channel_dim=None, space_dim=None)[source]
Parameters:

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

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

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

  • event_dim (int | None)

  • channel_dim (int | None)

  • space_dim (int | None)

Methods

process_and_flatten(x)

Process array and flatten event dimensions.

process_event(batched_shape)

Separate batch and event dimensions from a complete shape.

Attributes

channel_axes

Axis indices for channel dimensions.

channel_dim

channel_shape

channel_size

Total number of elements in channel dimensions.

event_axes

Axis indices for event dimensions.

event_dim

event_shape

event_size

Total number of elements in event shape.

space_axes

Axis indices for spatial dimensions.

space_dim

space_shape

space_size

Total number of elements in spatial dimensions.
event_shape: tuple[int, ...] | None = None
space_shape: tuple[int, ...] | None = None
channel_shape: tuple[int, ...] | None = None
event_dim: int | None = None
space_dim: int | None = None
channel_dim: int | None = None
process_event(batched_shape)[source]

Separate batch and event dimensions from a complete shape.

Also infers event, channel, space shapes, if only dimensions were known.

Parameters:

batched_shape (tuple[int, ...]) – Complete shape including batch and event dimensions.

Returns:

Tuple of (batch_shape, event_shape_info) where event_shape_info is a new ShapeInfo instance with inferred event dimensions.

Raises:

  • RuntimeError – If event dimension is unknown.

  • AssertionError – If provided event_shape doesn’t match inferred shape.

process_and_flatten(x)[source]

Process array and flatten event dimensions.

Parameters:

x (Array) – Input array with batch and event dimensions.

Returns:

Tuple of (flattened_array, shape_info) where event dimensions are flattened into a single dimension.

property event_axes: tuple[int, ...]

Axis indices for event dimensions.

property channel_axes: tuple[int, ...]

Axis indices for channel dimensions.

property space_axes: tuple[int, ...]

Axis indices for spatial dimensions.

property event_size: int

Total number of elements in event shape.

property space_size: int

Total number of elements in spatial dimensions.

property channel_size: int

Total number of elements in channel dimensions.