Ctrl+K

render_sharded_shape

render_sharded_shape#

penzai.treescope.arrayviz.arrayviz.render_sharded_shape(sharding: jax.sharding.Sharding, shape_or_namedarray_struct: jax.ShapeDtypeStruct | named_axes.NamedArrayBase | tuple[int, ...] | Any, rows: list[int | named_axes.AxisName] | None = None, columns: list[int | named_axes.AxisName] | None = None) → ArrayvizRendering[source]#

Renders the sharding an array would have, based on its shape.

Parameters:

  • sharding – A sharding to visualize.

  • shape_or_namedarray_struct – Either an arbitrary object with a .shape attribute, a tuple of integers, or a NamedArray wrapping a jax.lax.ShapeDtypeStruct.

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

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

Returns:

A rendering of the result of sharding an array with this shape using the given sharding.

Contents