jax.debug.visualize_sharding#

jax.debug.visualize_sharding(shape, sharding, *, use_color=True, scale=1.0, min_width=9, max_width=80, color_map=None)[source]#

Visualizes a Sharding using rich.

Parameters