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 class

CompiledCallable

CompiledCallable

class max.experimental.compilation.CompiledCallable(staged, weights=None, *, session=None)

source

Bases: Generic[_P, _R]

A compiled function that runs on tensors.

compile() returns one. Call it with the same arguments as the original function, passing a Tensor for each input type. The weights and the signal buffers are bound when the object is created, so the first call has no extra setup cost.

The following example compiles a function, exports it, and runs it:

from max.driver import CPU
from max.dtype import DType
from max.experimental import compilation
from max.experimental.sharding import TensorLayout
from max.experimental.tensor import Tensor

def scale(x: Tensor) -> Tensor:
    return x * 2

x_type = TensorLayout(DType.float32, [3], CPU())

run = compilation.compile(scale)(x_type)
run.export_mef("scale.mef")
out = run(Tensor.ones([3], device=CPU()))

Parameters:

• staged (StagedGraph[_P, _R]) – The traced graph to compile.
• weights (Mapping[str, DLPackArray] | None) – The weight registry: a mapping from the name of each weight the graph declares to its data, with one entry per shard for a distributed weight.
• session (InferenceSession | None) – The session that compiles and initializes the graph, including its precompiled or exported MEFs and its kernel options. Defaults to a session shared by the whole process.

engine_model

property engine_model: Model

source

Returns the underlying Model.

Use it to drive the model directly, for example to capture and replay it:

run.engine_model.capture(key, *buffers, *run.signal_buffers)
run.engine_model.replay(key, *buffers, *run.signal_buffers)

Raises:

RuntimeError – If the callable has no initialized model, as a deserialized one does not.

execute_raw()

execute_raw(*buffers)

source

Runs the compiled graph on buffers, skipping the argument checks.

Parameters:

buffers (Buffer) – One buffer per graph input, in graph order. A distributed argument takes one buffer per shard. Do not pass the signal buffers; they are added automatically, after the other inputs or where signal_buffer_index placed them.

Returns:

One buffer per graph output, in graph order.

Return type:

list[Buffer]

export_mef()

export_mef(path)

source

Writes the compiled graph to a MEF file.

Load the file with max.engine.read() to skip compilation the next time.

Parameters:

path (str | Path) – The path of the file to write.

Return type:

None

from_compiled()

classmethod from_compiled(artifact, engine_model, in_layouts, in_tree, out_layouts, out_tree, weights=None)

source

Wraps a graph compiled elsewhere, given its signature.

Use this for a graph that was traced without stage(), such as a Module graph. Any signal buffers are among its inputs.

Parameters:

• artifact (CompiledModel) – The compiled artifact, holding the one graph.
• engine_model (Model | None) – The model initialized from artifact, or None if the callable only describes the graph.
• in_layouts (Sequence[TensorLayout]) – The layout of each tensor argument, in graph order.
• in_tree (TreeDef) – The tree structure of the arguments.
• out_layouts (Sequence[TensorLayout]) – The layout of each tensor result, in graph order.
• out_tree (TreeDef) – The tree structure of the return value.
• weights (Mapping[str, DLPackArray] | None) – The weight registry engine_model was initialized with.

Returns:

The callable that runs engine_model.

Return type:

CompiledCallable[…, Any]

from_engine_model()

classmethod from_engine_model(staged, artifact, engine_model, weights=None)

source

Wraps a graph that was compiled and initialized elsewhere.

Use this when the caller compiles and initializes staged.graph itself, for example together with other graphs.

Parameters:

• staged (StagedGraph[~_P, _R]) – The traced graph that artifact was compiled from.
• artifact (CompiledModel) – The compiled artifact.
• engine_model (Model) – The model initialized from artifact.
• weights (Mapping[str, DLPackArray] | None) – The weight registry engine_model was initialized with.

Returns:

The callable that runs engine_model.

Return type:

CompiledCallable[~_P, _R]

in_layouts

in_layouts: list[TensorLayout]

source

See StagedGraph.in_layouts.

in_tree

in_tree: TreeDef

source

See StagedGraph.in_tree.

out_layouts

out_layouts: list[TensorLayout]

source

See StagedGraph.out_layouts.

out_tree

out_tree: TreeDef

source

See StagedGraph.out_tree.

signal_buffers

property signal_buffers: list[Buffer]

source

Returns the signal buffers for multi-device collectives.

The list is empty when the graph runs on a single device. The buffers are allocated once per set of devices and shared by every compiled callable that runs on them.

weights

property weights: Mapping[str, DLPackArray]

source

Returns the weight registry this callable was compiled with.

The mapping is read-only. To run the same graph with different weights, create a new CompiledCallable from the same StagedGraph.