faster_array_repr

penzai.treescope.ndarray_summarization.faster_array_repr(array: np.ndarray | jax.Array) → str

Computes repr(array), only copying the rendered array elements.

repr(array) on a very large jax Array can be slow, because it copies the entire array to host memory even when only a few elements are actually needed. We can avoid this by truncating the array on device before fetching it.

Parameters:

array – The array to summarize.

Returns:

A string representation of the array. May differ slightly from the ordinary repr, but should contain the same elements.