Skip to main content

struct

Graph

Represents a single MAX graph.

A Graph is a callable routine in MAX Engine, similar to a function in Mojo. Like functions, graphs have a name and signature. Unlike a function, which follows an imperative programming model, a Graph follows a dataflow programming model, using lazily-executed, parallel operations instead of sequential instructions.

When you instantiate a graph, you must specify the input shapes as one or more TensorType or ListType values. Then, build a sequence of ops and set the graph output with output(). For example:

from max.graph import Type, Graph, TensorType, ops
from tensor import Tensor, TensorShape

def build_model() -> Graph:
var graph = Graph(TensorType(DType.float32, 2, 6))

var matmul_constant_value = Tensor[DType.float32](TensorShape(6, 1), 0.15)
var matmul_constant = graph.constant(matmul_constant_value)

var matmul = graph[0] @ matmul_constant
var relu = ops.elementwise.relu(matmul)
var softmax = ops.softmax(relu)
graph.output(softmax)

return graph

You can't call a Graph directly from Mojo. You must compile it and execute it with MAX Engine. For more detail, see the tutorial about how to build a graph with MAX Graph.

Implemented traits

AnyType, CollectionElement, Copyable, Movable, Stringable

Methods

__init__

__init__(inout self: Self, /, in_type: Type)

Constructs a new Graph with a single input type.

Although a Graph is technically valid once constructed, it is not usable for inference until you specify outputs by calling output(). Check the graph validity by calling verify().

Args:

__init__(inout self: Self, /, in_types: List[Type], out_types: List[Type] = List())

Constructs a new Graph using the default graph name.

Although a Graph is technically valid once constructed, it is not usable for inference until you specify outputs by calling output(). Check the graph validity by calling verify().

Args:

  • in_types (List[Type]): The graph's input types, as one or more TensorType or ListType values.
  • out_types (List[Type]): The graph's output types, as one or more TensorType or ListType values. Deprecated. This will be inferred by the output call.

__init__(inout self: Self, /, name: String, in_types: List[Type], out_types: List[Type] = List())

Constructs a new Graph with a custom graph name.

Although a Graph is technically valid once constructed, it is not usable for inference until you specify outputs by calling output(). Check the graph validity by calling verify().

Args:

  • name (String): A name for the graph.
  • in_types (List[Type]): The graph's input types, as one or more TensorType or ListType values.
  • out_types (List[Type]): The graph's output types, as one or more TensorType or ListType values. Deprecated. This will be inferred by the output call.

__getitem__

__getitem__(self: Self, n: Int) -> Symbol

Returns the n'th argument of this Graph.

By argument, we mean the graph input. For example, graph[0] gets the first input and graph[1] gets the second input (as specified with the Graph constructor's in_types).

This provides the argument as a Symbol, which you can use as input to other nodes in the graph.

Args:

  • n (Int): The argument number. First argument is at position 0.

Returns:

A Symbol representing the argumen't symbolic value, as seen from within the Graph's body.

Raises:

If n is not a valid argument number.

__str__

__str__(self: Self) -> String

Returns a String representation of this Graph.

The representation uses a MLIR textual format. The format is subject to change and should only be used for debugging pruposes.

Returns:

A human-readable string representation of the graph.

verify

verify(self: Self)

Verifies the Graph and its contents.

Examples of cases when a Graph may not be valid (the list is not exhaustive):

  1. it has an output op whose types don't match its out_types
  2. it has an op with an invalid name, number, type of operands, output types, etc.
  3. it contains cycles

Raises:

If the Graph did not pass verification. In this case it will also print a diagnostic message indicating the error.

nvop

nvop(self: Self, name: String, inputs: List[Symbol] = List(), out_types: List[Type] = List(), attrs: List[NamedAttribute] = List(), enable_result_type_inference: Bool = false) -> List[Symbol]

Adds a new node to the Graph.

The node represents a single MAX Graph operation.

This is a very low level API meant to enable creating any supported op. In general, it's less ergonomic compared to the higher level helpers in the ops package.

Note that these nodes don't take concrete values as inputs, but rather symbolic values representing the outputs of other nodes or the Graphs arguments.

Args:

  • name (String): The name of the operation to use.
  • inputs (List[Symbol]): The list of symbolic operands.
  • out_types (List[Type]): The list of output types.
  • attrs (List[NamedAttribute]): Any attributes that the operation might require.
  • enable_result_type_inference (Bool): Will infer the result type if True.

Returns:

The symbolic outputs of the newly-added node.

op

op(self: Self, name: String, out_type: Type, attrs: List[NamedAttribute] = List()) -> Symbol

Adds a new single-output, nullary node to the Graph.

See Graph.nvop for details. This overload can be used for operations that take no inputs and return a single result, such as mo.constant.

Args:

  • name (String): The name of the operation to use.
  • out_type (Type): The output types.
  • attrs (List[NamedAttribute]): Any attributes that the operation might require.

Returns:

The symbolic output of the newly-added node.

op(self: Self, name: String, inputs: List[Symbol], out_type: Type, attrs: List[NamedAttribute] = List()) -> Symbol

Adds a new single-output node to the Graph.

See Graph.nvop for details. This overload can be used for operations that return a single result.

Args:

  • name (String): The name of the operation to use.
  • inputs (List[Symbol]): The list of symbolic operands.
  • out_type (Type): The output types.
  • attrs (List[NamedAttribute]): Any attributes that the operation might require.

Returns:

The symbolic output of the newly-added node.

constant

constant[dtype: DType](self: Self, owned value: Tensor[dtype]) -> Symbol

Adds a node representing a mo.constant operation.

The value of this constant will have the type TensorType with the same shape and dtype as value.

Parameters:

  • dtype (DType): The constant tensor's element type.

Args:

  • value (Tensor[dtype]): The constant's value.

Returns:

The symbolic output of this node.

vector

vector[dtype: DType](self: Self, values: List[SIMD[dtype, 1]]) -> Symbol

Adds a node representing a mo.constant operation.

The value of this constant will have the type TensorType with 1-D shape, consistent with the size of values.

Parameters:

  • dtype (DType): The constant tensor's element type.

Args:

  • values (List[SIMD[dtype, 1]]): A vector represneting the constant's value.

Returns:

The symbolic output of this node.

scalar

scalar[dtype: DType](self: Self, value: SIMD[dtype, 1], rank: Int = 0) -> Symbol

Adds a node representing a mo.constant operation.

The value of this constant will have the type scalar TensorType (0-D shape), when rank is 0, or a higher-rank TensorType of a single element.

Parameters:

  • dtype (DType): The constant tensor's element type.

Args:

  • value (SIMD[dtype, 1]): The constant's value.
  • rank (Int): The output tensor's rank.

Returns:

The symbolic output of this node.

scalar(self: Self, value: Int, dtype: DType) -> Symbol

Adds a node representing a mo.constant operation.

The value of this constant will have the type TensorType of the same element type as dtype, and scalar (0-D) shape.

Args:

  • value (Int): The scalar value.
  • dtype (DType): The constant's element type.

Returns:

The symbolic output of this node.

Raises:

If value cannot be instantiated as a tensor of element dtype.

scalar(self: Self, value: SIMD[float64, 1], dtype: DType) -> Symbol

Adds a node representing a mo.constant operation.

The value of this constant will have the type TensorType of the same element type as dtype, and scalar (0-D) shape.

Args:

  • value (SIMD[float64, 1]): The scalar value.
  • dtype (DType): The constant's element type.

Returns:

The symbolic output of this node.

Raises:

If value cannot be instantiated as a tensor of element dtype.

range

range[dtype: DType](self: Self, start: SIMD[dtype, 1], stop: SIMD[dtype, 1], step: SIMD[dtype, 1]) -> Symbol

Adds a node representing a mo.range operation.

Parameters:

  • dtype (DType): The output tensor's element type.

Args:

  • start (SIMD[dtype, 1]): The starting value.
  • stop (SIMD[dtype, 1]): The end value (exclusive).
  • step (SIMD[dtype, 1]): The step value.

Returns:

The symbolic output of this node.

full

full[dtype: DType](self: Self, value: SIMD[dtype, 1], dims: List[Symbol]) -> Symbol

Creates a constant-valued symbolic tensor of a specified shape.

Parameters:

  • dtype (DType): The output tensor's element type.

Args:

  • value (SIMD[dtype, 1]): The value to fill the resulting tensor with.
  • dims (List[Symbol]): The shape dimensions of the zero-valued tensor.

Returns:

A symbolic tensor of the specified shape and dtype, where every value is the specified fill value.

output

output(inout self: Self, outputs: List[Symbol])

Adds an output for the graph.

This is a special node that all graphs must have in order to deliver inference results. The outs symbol given here must match the shape and type of the out_types given when constructing the graph.

Args:

  • outputs (List[Symbol]): The return values, usually the result from one or more ops.