infer_balanced_truncation

penzai.treescope.ndarray_summarization.infer_balanced_truncation(shape: tuple[int, ...], maximum_size: int, cutoff_size_per_axis: int, minimum_edge_items: int, doubling_bonus: float = 10.0) → tuple[int | None, ...]

Infers a balanced truncation from a shape.

This function computes a set of truncation sizes for each axis of the array such that it obeys the constraints about array and axis sizes, while also keeping the relative proportions of the array consistent (e.g. we keep more elements along axes that were originally longer). This means that the aspect ratio of the truncated array will still resemble the aspect ratio of the original array.

To avoid very-unbalanced renderings and truncate longer axes more than short ones, this function truncates based on the square-root of the axis size by default.

Parameters:

• shape – The shape of the array we are truncating.

• maximum_size – Maximum number of elements of an array to show. Arrays larger than this will be truncated along one or more axes.

• cutoff_size_per_axis – Maximum number of elements of each individual axis to show without truncation. Any axis longer than this will be truncated, with their visual size increasing logarithmically with the true axis size beyond this point.

• minimum_edge_items – How many values to keep along each axis for truncated arrays. We may keep more than this up to the budget of maximum_size.

• doubling_bonus – Number of elements to add to each axis each time it doubles beyond cutoff_size_per_axis. Used to make longer axes appear visually longer while still keeping them a reasonable size.

Returns:

A tuple of edge sizes. Each element corresponds to an axis in shape, and is either None (for no truncation) or an integer (corresponding to the number of elements to keep at the beginning and and at the end).