Large arrays of PRNG keys cannot be deferred

Hi! First, thanks for this amazing tool! I noticed that your tool fails to render large PRNG key arrays, see: <img width="1213" height="609" alt="Image" src="https://github.com/user-attachments/assets/b1c5005d-95c9-4187-a8c5-027ae7fb7675" /> Clicking on the arrow reveals the following traceback: ``` <TypeError during deferred rendering Traceback (most recent call last): File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/treescope/lowering.py", line 381, in _render_to_html_as_root_streaming replacement_part = deferred.thunk(layout_decision) File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/treescope/external/jax_support.py", line 613, in _thunk node_repr = faster_array_repr(node) File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/treescope/external/jax_support.py", line 285, in faster_array_repr array_edges, _ = truncate_array_and_mask( ~~~~~~~~~~~~~~~~~~~~~~~^ array, ^^^^^^ np.ones((1,) * array.ndim, dtype=jnp.bool_), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ edge_items_per_axis=tuple(edge_items_per_axis), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ) ^ File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/treescope/external/jax_support.py", line 251, in truncate_array_and_mask return fn(array, mask, (), edge_items_per_axis) File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/treescope/external/jax_support.py", line 139, in _truncate_part_with_slices array = xnp.array(array) File "/export/home/eertmans/repositories/sampling-spirals/.venv/lib/python3.13/site-packages/jax/_src/prng.py", line 305, in __array__ raise TypeError("JAX array with PRNGKey dtype cannot be converted to a NumPy array." " Use jax.random.key_data(arr) if you wish to extract the underlying" " integer array.") TypeError: JAX array with PRNGKey dtype cannot be converted to a NumPy array. Use jax.random.key_data(arr) if you wish to extract the underlying integer array. > ``` I think the error is similar to https://github.com/google-deepmind/treescope/issues/52, but here applies to generating the summary of an array, which does not occur for smaller arrays: <img width="1213" height="691" alt="Image" src="https://github.com/user-attachments/assets/012c151b-311b-4941-8372-1f90c4fead42" /> *Notes:* I am using `treescope==0.1.10` and `jax==0.7.1`.