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)

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]]]) –