Array visualization fails for large, sharded arrays

Hi, When I try to visualize sharded arrays above a certain size like this: ``` mesh = jax.sharding.Mesh(mesh_utils.create_device_mesh((1,4)), axis_names=("dp", "tp"), axis_types=(jax.sharding.AxisType.Explicit,jax.sharding.AxisType.Explicit)) jax.device_put(jnp.zeros((5120, 5120)), jax.sharding.NamedSharding(mesh, jax.P("tp", None))) ``` I get the following error: ``` <ShardingTypeError during deferred rendering Traceback (most recent call last): File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/lowering.py", line 381, in _render_to_html_as_root_streaming replacement_part = deferred.thunk(layout_decision) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/_internal/api/array_autovisualizer.py", line 286, in _thunk return self._autovisualize_array( ^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/_internal/api/array_autovisualizer.py", line 130, in _autovisualize_array adapter.get_array_data_with_truncation( File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/external/jax_support.py", line 521, in get_array_data_with_truncation array, mask = truncate_array_and_mask(array, mask, edge_items_per_axis) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/external/jax_support.py", line 252, in truncate_array_and_mask return fn(array, mask, (), edge_items_per_axis) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback return fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 260, in cache_miss _python_pjit_helper(fun, jit_info, *args, **kwargs) File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 137, in _python_pjit_helper p, args_flat = _infer_params(fun, jit_info, args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 629, in _infer_params return _infer_params_internal(fun, ji, args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 653, in _infer_params_internal p, args_flat = _infer_params_impl( ^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 552, in _infer_params_impl jaxpr, consts, jaxpr_for_top_jit, out_avals = _create_pjit_jaxpr( ^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/linear_util.py", line 472, in memoized_fun ans = call(fun, *args) ^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1350, in _create_pjit_jaxpr jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(fun, in_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 354, in wrapper return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/interpreters/partial_eval.py", line 2337, in trace_to_jaxpr_dynamic ans = fun.call_wrapped(*in_tracers) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/linear_util.py", line 212, in call_wrapped return self.f_transformed(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 73, in flatten_fun ans = f(*py_args, **py_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 314, in _argnames_partial return _fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 288, in _argnums_partial return _fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/linear_util.py", line 397, in _get_result_paths_thunk ans = _fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/external/jax_support.py", line 172, in _truncate_part_with_slices result_a, valid_a = _truncate_part_with_slices( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/external/jax_support.py", line 172, in _truncate_part_with_slices result_a, valid_a = _truncate_part_with_slices( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/treescope/external/jax_support.py", line 146, in _truncate_part_with_slices truncated_array = array[prefix_slices] ~~~~~^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 1125, in op return getattr(self.aval, f"_{name}")(self, *args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 660, in _getitem return indexing.rewriting_take(self, item) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/numpy/indexing.py", line 639, in rewriting_take result = _attempt_rewriting_take_via_slice(arr, idx, mode, out_sharding) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/numpy/indexing.py", line 602, in _attempt_rewriting_take_via_slice arr = slicing.slice( ^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/slicing.py", line 111, in slice return slice_p.bind(operand, start_indices=tuple(start_indices), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 576, in bind return self._true_bind(*args, **params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 592, in _true_bind return self.bind_with_trace(prev_trace, args, params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 602, in bind_with_trace return trace.process_primitive(self, args, params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/interpreters/partial_eval.py", line 2062, in process_primitive return self.default_process_primitive( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/interpreters/partial_eval.py", line 2076, in default_process_primitive out_avals, effs = _cached_abstract_eval(primitive, *aval_qdds, **params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/util.py", line 298, in wrapper return cached(config.trace_context() if trace_context_in_key else _ignore(), ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/util.py", line 292, in cached return f(*args, **kwargs) ^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/interpreters/partial_eval.py", line 1900, in _cached_abstract_eval return primitive.abstract_eval(*aval_qdds, **params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 641, in abstract_eval_ return abstract_eval(*args, **kwargs), no_effects ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/utils.py", line 139, in standard_abstract_eval out_shape, out_dtype, out_sharding = call_shape_dtype_sharding_rule( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/utils.py", line 105, in call_shape_dtype_sharding_rule out_shardings = call_sharding_rule( ^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/utils.py", line 94, in call_sharding_rule out_sharding = sh_rule(*avals, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/slicing.py", line 1418, in _slice_sharding_rule return _get_sharding_for_varying_out_shape(out_shape, operand, 'slicing') ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/user/ml-scratch/.venv/lib/python3.12/site-packages/jax/_src/lax/slicing.py", line 1398, in _get_sharding_for_varying_out_shape raise core.ShardingTypeError( jax._src.core.ShardingTypeError: slicing on sharded dims where out dim (31) is not divisible by mesh axes (4) with spec (tp) is not implemented. ``` I suspect it's because https://github.com/google-deepmind/treescope/blob/main/treescope/_internal/arrayviz_impl.py#L432 does not appear to account for sharding constraints like this. Thanks.