summarization_threshold

penzai.treescope.ndarray_summarization.summarization_threshold: context.ContextualValue[Mapping[str, int | None]]

Threshold for summarization of NDArrays for each backend.

This threshold determines the largest number of elements we will summarize with summary statistics (e.g. mean, standard deviation) when rendering in treescope. Larger values may make it slower to display large NDArrays.

Each key should be the name of a JAX array platform, e.g. “cpu” or “tpu”. It can also be “numpy” to refer to Numpy arrays, or “default” to refer to any other accelerator. The value is the size of the array at which point we avoid showing summary statistics. None means no limit.

This configuration argument is intended to be set at the top level by the user, e.g. in IPython.