For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

visualize

`visualize()`

max.experimental.sharding.visualize(value, *, title=None, blocks=True, tree=True)

source

Draws a mesh, or a tensor’s slices, as blocks beside their tree.

The blocks on the left give one cell per device (for a mesh) or one cell per slice, sized in proportion to the local shape; the tree on the right gives one level per mesh axis and the slice held at every node. Both describe the same value, so they share one header.

A dict draws its values side by side, smaller, each under its key, so that related values such as an op’s inputs and output read in one row:

from max.driver import CPU
from max.dtype import DType
from max.experimental.sharding import (
    DeviceMapping,
    DeviceMesh,
    Partial,
    Sharded,
    TensorLayout,
    visualize,
)

mesh = DeviceMesh(tuple(CPU() for _ in range(4)), (4,), ("tp",))

def layout(shape, placement):
    return TensorLayout(
        DType.float32, shape, DeviceMapping(mesh, (placement,))
    )

lhs = layout([4, 8], Sharded(1))
rhs = layout([8, 8], Sharded(0))
out = layout([4, 8], Partial())
visualize({"lhs": lhs, "rhs": rhs, "out": out}, tree=False)

Parameters:

value (Any) – A DeviceMesh, a TensorLayout or a tensor, or a dict of them keyed by title.
title (str | None) – Text prepended to the header, such as the variable name. For a dict, a line above the whole row.
blocks (bool) – Whether to draw the blocks.
tree (bool) – Whether to draw the tree. The blocks cannot show a mesh of more than two axes or more than two sharded tensor axes, so such a value is drawn as the tree alone however these two are set.

Raises:

TypeError – If value is none of those.
ValueError – If both blocks and tree are false, or if a dict is empty.

Return type:

None

visualize()​

`visualize()`