google-deepmind / google-deepmind/treescope
Array visualization fails for large, sharded arrays
Open
- Dominant language
- Python
- Stars
- 474
- Forks
- 28
- PR merge metrics
- No merged PRs in 30d
Description
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:
```
Contributor guide
Assessment
This issue has not been assessed yet.