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:
- in_type (
Type
): The graph's input type, as a singleTensorType
orListType
value.
__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 moreTensorType
orListType
values. - out_types (
List[Type]
): The graph's output types, as one or moreTensorType
orListType
values. Deprecated. This will be inferred by theoutput
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 moreTensorType
orListType
values. - out_types (
List[Type]
): The graph's output types, as one or moreTensorType
orListType
values. Deprecated. This will be inferred by theoutput
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):
- it has an
output
op whose types don't match itsout_types
- it has an op with an invalid name, number, type of operands, output types, etc.
- 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
Graph
s 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.