Ctrl+K

render_array_sharding

render_array_sharding#

penzai.treescope.arrayviz.arrayviz.render_array_sharding(array: jax.Array | named_axes.NamedArray, rows: list[int | named_axes.AxisName] | None = None, columns: list[int | named_axes.AxisName] | None = None) → ArrayvizRendering[source]#

Renders the sharding of an array.

Parameters:

  • array – The array whose sharding we should render.

  • rows – Optional explicit ordering of axes for the visualization rows.

  • columns – Optional explicit ordering of axes for the visualization columns.

Returns:

A rendering of that array’s sharding.

Contents