Mojo 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 max.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
,
Formattable
,
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, 0], out_types: List[Type, 0] = 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, 0]
): The graph's input types, as one or moreTensorType
orListType
values. - βout_types (
List[Type, 0]
): 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, 0], out_types: List[Type, 0] = 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, 0]
): The graph's input types, as one or moreTensorType
orListType
values. - βout_types (
List[Type, 0]
): The graph's output types, as one or moreTensorType
orListType
values. Deprecated. This will be inferred by theoutput
call.
__init__(inout self: Self, path: Path)
Constructs a new Graph
from a MLIR file.
Experimental. Recreates a graph from an MLIR file.
Args:
- βpath (
Path
): The path of the MLIR file.
__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.
debug_str
β
debug_str(self: Self, pretty_print: Bool = 0) -> String
__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.
format_to
β
format_to(self: Self, inout writer: Formatter)
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.
layer
β
layer(inout self: Self, name: String) -> _GraphLayerContext
Creates a context manager for a graph layer.
Graph layers don't have a functional meaning for graph execution. They help provide debug and visualization information, tagging nodes in the graph with informal information about the structure of the overall graph.
Args:
- βname (
String
): The name of the layer.
Returns:
A context manager. Inside this context, the layer will be "active".
current_layer
β
current_layer(self: Self) -> String
Returns the full path of the current layer.
This is a .
-separated string of each nested layer context created
by Graph.layer()
.
Returns:
The full path of the current layer.
nvop
β
nvop(self: Self, name: String, inputs: List[Symbol, 0] = List(), out_types: List[Type, 0] = List(), attrs: List[NamedAttribute, 0] = List(), enable_result_type_inference: Bool = 0) -> List[Symbol, 0]
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, 0]
): The list of symbolic operands. - βout_types (
List[Type, 0]
): The list of output types. - βattrs (
List[NamedAttribute, 0]
): 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, 0] = 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, 0]
): Any attributes that the operation might require.
Returns:
The symbolic output of the newly-added node.
op(self: Self, name: String, inputs: List[Symbol, 0], out_type: Type, attrs: List[NamedAttribute, 0] = 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, 0]
): The list of symbolic operands. - βout_type (
Type
): The output types. - βattrs (
List[NamedAttribute, 0]
): Any attributes that the operation might require.
Returns:
The symbolic output of the newly-added node.
op(self: Self, name: String, inputs: List[Symbol, 0], attrs: List[NamedAttribute, 0] = List()) -> Symbol
Adds a new single-output node to the Graph
with result type inference.
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, 0]
): The list of symbolic operands. - βattrs (
List[NamedAttribute, 0]
): 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.
quantize
β
quantize[encoding: QuantizationEncoding](self: Self, owned value: Tensor[float32]) -> Symbol
Quantizes a tensor using a specific quantization encoding.
This takes the full-precision value
as owned data and frees it.
The resulting quantized constant is allocated and owns its data.
Parameters:
- βencoding (
QuantizationEncoding
): Describes a specific quantization encoding such as Q4_0.
Args:
- βvalue (
Tensor[float32]
): Full-precision value to quantize.
Returns:
Symbol representing the quantized constant.
vector
β
vector[dtype: DType](self: Self, values: List[SIMD[dtype, 1], 0]) -> 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], 0]
): 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
Creates a sequence of numbers. The sequence goes from start
with increments of size step
up to (but not including) stop
. All arguments are mandatory and must have the same element type.
Note the following restrictions on input values:
step
must be non-zerostop - start
must be zero or have the same sign asstep
Parameters:
- βdtype (
DType
): The output tensor's element type.
Args:
- βstart (
SIMD[dtype, 1]
): The start of the range to generate. - βstop (
SIMD[dtype, 1]
): The range will be generated up to, but not including, this value. - βstep (
SIMD[dtype, 1]
): The step size for the range.
Returns:
A symbolic tensor value containing the defined range of values.
range(self: Self, start: Symbol, stop: Symbol, step: Symbol, out_dim: Dim) -> Symbol
Creates a sequence of numbers. The sequence goes from start
with increments of size step
up to (but not including) stop
. All arguments are mandatory and must have the same element type.
Note the following restrictions on input values:
step
must be non-zerostop - start
must be zero or have the same sign asstep
Args:
- βstart (
Symbol
): The start of the range to generate. - βstop (
Symbol
): The range will be generated up to, but not including, this value. - βstep (
Symbol
): The step size for the range. - βout_dim (
Dim
): The expected output dimensions returned by the range op. These will be assert at graph execution time to be correct.
Returns:
A symbolic tensor value containing the defined range of values.
full
β
full[dtype: DType](self: Self, value: SIMD[dtype, 1], *dims: Dim) -> 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 (
Dim
): 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.
full[dtype: DType](self: Self, value: SIMD[dtype, 1], dims: List[Dim, 0], location: Optional[_SourceLocation] = #kgen.none) -> 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[Dim, 0]
): The shape dimensions of the zero-valued tensor. - βlocation (
Optional[_SourceLocation]
): An optional location for a more specific error message.
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, 0])
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, 0]
): The return values, usually the result from one or more ops.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!
If you'd like to share more information, please report an issue on GitHub
π What went wrong?