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)

使用 rich 可视化 Sharding。

参数：

• shape (Sequence[int])

• sharding (Sharding)

• use_color (bool)

• scale (float)

• min_width (int)

• max_width (int)

• color_map (ColorMap | None | None)