mr2.utils.broadcast_shapes_except

mr2.utils.broadcast_shapes_except(shapes: Sequence[Sequence[int]], dim: int) → list[tuple[int, ...]][source]

Broadcast shapes against each other, except along one dimension.

Inputs are first left-padded with ones to a common number of dimensions (matching PyTorch right-aligned broadcasting). All dimensions except dim are then broadcast normally. Each output shape keeps its input’s original size along dim.

Parameters:

• shapes (Sequence[Sequence[int]]) – Shapes to broadcast.

• dim (int) – Dimension to exclude from broadcasting. Negative values are normalized relative to the common (left-padded) number of dimensions.

Returns:

A list of per-input target shapes that agree in every dimension except dim.