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.tensor.sharding.visualize(value, *, title=None, blocks=True, tree=True)

source

Draws a mesh or a tensor’s slices as blocks beside the tree they share.

The blocks on the left give one cell per device for a mesh, and one cell per slice for a tensor, with the devices that share a slice listed together. Cells are sized in proportion to the local shape when every dimension is static, and equally otherwise. A Partial or Unknown cell is faded and labelled: what each device holds is not a piece of the drawn tensor. The tree gives one level per mesh axis and the value 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.tensor.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 (DeviceMesh | TensorLayout | _HasLayout | Mapping[str, DeviceMesh | TensorLayout | _HasLayout]) – A DeviceMesh, a TensorLayout, or anything with a layout property that is a TensorLayout (such as a Tensor), or a dict of them keyed by title. A bare DeviceMapping is not accepted; wrap it in a TensorLayout to give it a dtype and a shape.
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()`