infer_rows_and_columns

penzai.treescope.arrayviz.arrayviz.infer_rows_and_columns(axis_sizes: dict[int | named_axes.AxisName, int], unassigned: Sequence[int | named_axes.AxisName] | None = None, known_rows: Sequence[int | named_axes.AxisName] = (), known_columns: Sequence[int | named_axes.AxisName] = ()) → tuple[list[int | named_axes.AxisName], list[int | named_axes.AxisName]]

Infers an ordered assignment of axis indices or names to rows and columns.

The unassigned axes are sorted by size and then assigned to rows and columns to try to balance the total number of elements along the row and column axes. Curently uses a greedy algorithm with an adjustment to try to keep columns longer than rows, except when there are exactly two axes and both are positional, in which case it lays out axis 0 as the rows and axis 1 as the columns.

Parameters:

• axis_sizes – Mapping from axis indices or names to their axis size.

• unassigned – Sequence of unassigned axis indices or names. Inferred from the axis_sizes if not provided.

• known_rows – Sequence of axis indices or names that must map to rows.

• known_columns – Sequence of axis indices or names that must map to columns.

Returns:

Tuple (rows, columns) of assignments, which consist of known_rows and known_columns followed by the remaining unassigned axes in a balanced order.