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)))
<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.
Thanks.
Hi,
When I try to visualize sharded arrays above a certain size like this:
I get the following error:
I suspect it's because main/treescope/_internal/arrayviz_impl.py#L432 does not appear to account for sharding constraints like this.
Thanks.