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 (Callable[[float], tuple[float, float, float, float]] | None)
