Sharding Visualisation/Debugging

[Findus23]

Recently I have been working a lot with sharding in JAX and trying to wrap my head around how it can be implemented in a more complex existing codebasis.
I really like the visualisation of jax.debug.visualize_array_sharding(), but it has the major limitation that it only supports arrays with at most two dimensions. As in my case every array has a shape of at least 3 dimensions, I can't use it.

So I wrote a fork of this function that works around this by ignoring non-sharded dimensions until a 2D visualisation is possible.

array = jax.numpy.zeros(shape=(N, N, N), dtype=jax.numpy.float32)
array = jax.device_put(array, NamedSharding(mesh, P(None, "gpus")))
sharding_vis(array)

─────────── showing dims [0, 1] from original shape (128, 128, 128) ────────────
┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
│ CPU 0 │ CPU 1 │ CPU 2 │ CPU 3 │ CPU 4 │ CPU 5 │ CPU 6 │ CPU 7 │
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
│       │       │       │       │       │       │       │       │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘

And while thinking about it, I realized that in higher dimensions a simple text representation feels more intuitive to me than the table-view. Therefore I created another function, that works similar, but instead shows general information about the array (shape, dtype, nbytes), but also supports all sharding types and prints the axis along which the array is sharded.

array = jax.numpy.zeros(shape=(N, N, N), dtype=jax.numpy.float32)
array = jax.device_put(array, NamedSharding(mesh, P(None, "gpus")))
sharding_info(array, "some_array")

╭────────────────── some_array ───────────────────╮
│ shape: (128, 128, 128)                          │
│ dtype: float32                                  │
│ size: 8.0 MiB                                   │
│ NamedSharding: P(None, 'gpus')                  │
│ axis 1 is sharded: CPU 0 contains 0:16 (of 128) │
╰─────────────────────────────────────────────────╯

Keep in mind that the code is a bit experimental and might include some sharding misconceptions of mine, but it already helped me a lot to understand how the array sharding propagates through the calculations.
https://github.com/Findus23/jax-array-info/

You can find a few examples of the possible output here:
https://github.com/Findus23/jax-array-info/blob/main/tests/jaxtest.py