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 (布尔值)

• scale (浮点数)

• min_width (整数)

• max_width (整数)

• color_map (ColorMap | None | None)