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)

source

Traces and compiles fn.

Call the returned function with one spec per tensor argument of fn to get the CompiledCallable; call that on real tensors.

Tensor arguments are given as specs. A TensorLayout is a layout, and an argument given one is read-only; pass a BufferLayout for a buffer the callable may also store through. Only a layout may declare a boundary: a live Tensor is refused, because its dims are whatever it currently holds and taking them would fix every one. Pass tensor.layout to do that on purpose.

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_spec = TensorLayout(DType.float32, ["batch", 2], CPU())
w = Tensor.ones([2], device=CPU()) * 3

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

Parameters:

fn (Callable[[~_P], _R]) – The callable to compile, over Tensor values or containers of them.
weights (Mapping[str, DLPackArray] | None) – Data for the external constants the graph declares, keyed as the graph names them, one entry per shard of a distributed weight.
name (str | None) – The graph’s name. Defaults to fn’s own name.
custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries.
allow_subgraphs (bool) – Whether as_subgraph() bodies become shared subgraphs rather than inlining into the caller.
signal_devices (Iterable[Device]) – Devices taking part in collectives beyond what the specs span.
is_device_graph (bool) – Whether to record a device graph.

Returns:

A callable taking one spec per argument of fn and returning the CompiledCallable.

Return type:

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

compile()​

`compile()`