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

compile

`compile()`

max.experimental.compilation.compile(fn, *, weights=None, name=None, custom_extensions=(), allow_subgraphs=True, signal_devices=(), is_device_graph=False, session=None, signal_buffer_index=None)

source

Traces and compiles a function.

Returns a function that takes the input types of fn and returns a CompiledCallable. Call that on tensors to run fn.

A Tensor is not accepted as an input type, because its shape would fix every dimension of the graph to the tensor’s current size. Pass tensor.layout to do this deliberately.

The following example compiles a function that reads a weight, and supplies the weight’s value through weights:

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

w_type = TensorLayout(DType.float32, [2], CPU())

def layer(x: Tensor) -> Tensor:
    return x * F.constant_external("w", w_type)

x_type = TensorLayout(DType.float32, ["batch", 2], CPU())
w = Tensor.ones([2], device=CPU()) * 3

run = compilation.compile(layer, weights={"w": w})(x_type)
out = run(Tensor.ones([4, 2], device=CPU()))

Parameters:

fn (Callable[[~_P], _R]) – The function to compile. It takes and returns Tensor values, possibly inside lists, tuples or dictionaries.
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.
name (str | None) – The name of the graph. Defaults to the name of fn.
custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries.
allow_subgraphs (bool) – Whether calls made through as_subgraph() are traced as subgraphs. If False, they are inlined.
signal_devices (Iterable[Device]) – Additional accelerators to allocate signal buffers for, beyond the devices of the input types and of the tensors held by fn.
is_device_graph (bool) – Whether to build a device graph.
session (InferenceSession | None) – The session that compiles and initializes the graph. Defaults to a session shared by the whole process.
signal_buffer_index (int | None) – The graph input position of the signal buffers. See stage().

Returns:

A function that takes one input type per tensor argument of fn and returns the CompiledCallable.

Return type:

Callable[[…], CompiledCallable[~_P, _R]]

compile()​

`compile()`