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

GraphSpec

GraphSpec

class max.pipelines.lib.graph_specs.GraphSpec(module, name, input_types, weights=None, group=None, custom_extensions=(), is_device_graph=False, auto_cast=False, signal_buffer_index=None)

source

Bases: object

A graph a pipeline compiles, as the module that traces it.

A Module (ModuleV3) graph traces its forward through stage() and compiles on its own.

Parameters:

• module (Module | ModuleV3[..., Any]) – The module whose __call__ (forward for a ModuleV3 module) the graph traces.
• name (str) – The graph’s name. Graphs in one group need distinct names.
• input_types (Sequence[tree.Tree[InputType]] | Mapping[str, object]) – The graph’s input types: a sequence, or a named tree such as a module’s input_tree whose leaves are the inputs in order. A ModuleV3 graph takes a sequence, one type or tree of types per argument of its forward.
• weights (Mapping[str, DLPackArray] | None) – The weight values, keyed by name. Defaults to the module’s state_dict() once traced, or a ModuleV3 module’s own parameters.
• group (str | None) – Roles with the same group compile into one model, which holds each role’s graph as a named submodel. None compiles the role on its own. A ModuleV3 graph takes no group.
• custom_extensions (Sequence[Path]) – Paths to custom kernel libraries the graph uses.
• is_device_graph (bool) – Whether to build a device graph.
• auto_cast (bool) – Whether a ModuleV3 graph casts a float32 or bfloat16 weight to the dtype its parameter declares.
• signal_buffer_index (int | None) – Where a ModuleV3 graph takes its signal buffers among its inputs, to match an existing graph’s calling convention. None puts them after the last input, and only on a graph that runs on several devices.

auto_cast

auto_cast: bool = False

source

custom_extensions

custom_extensions: Sequence[Path] = ()

source

group

group: str | None = None

source

input_types

input_types: Sequence[tree.Tree[InputType]] | Mapping[str, object]

source

is_device_graph

is_device_graph: bool = False

source

module

module: Module | ModuleV3[..., Any]

source

name

name: str

source

signal_buffer_index

signal_buffer_index: int | None = None

source

trace()

trace(module=None)

source

Traces the graph, into module if given.

Raises:

ValueError – If module is given for a ModuleV3 graph.

Parameters:

module (Module | None)

Return type:

Graph

weights

weights: Mapping[str, DLPackArray] | None = None

source

weights_registry()

weights_registry()

source

Returns the weight values the graph binds, keyed by name.

Return type:

Mapping[str, DLPackArray]