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 shape (Sequence[int]) – sharding (Sharding) – use_color (bool) – scale (float) – min_width (int) – max_width (int) – color_map (Optional[Callable[[float], tuple[float, float, float, float]]]) –
