Mojo struct

Graph

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 main():
    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)

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

def main():
    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)

    print(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, UnknownDestructibility, Writable

Methods

`init`

__init__(out 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 single TensorType or ListType value.

__init__(out 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__(out 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.

__init__(out 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, 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, pretty_print: Bool = False) -> String

`str`

__str__(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.

`write_to`

write_to[W: Writer](self, mut writer: W)

`verify`

verify(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 its out_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(mut 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) -> 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, 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, 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, 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.

op(self, name: String, inputs: List[Symbol], attrs: List[NamedAttribute] = 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]): The list of symbolic operands.
attrs (List[NamedAttribute]): Any attributes that the operation might require.

Returns:

The symbolic output of the newly-added node.

`constant`

constant[dtype: DType](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, 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.

To quantize your model weights, follow these steps:

Import your trained model weights
Choose your quantization encodings
Apply quantization using Graph.quantize()

quantized_constant = graph.quantize[Q4_0Encoding](constant_value)
quantized_constant = graph.quantize[Q4_0Encoding](constant_value)

Use qmatmul() for matrix operations with quantized weights

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, 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, 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, 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, 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, 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-zero
stop - start must be zero or have the same sign as step

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, 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-zero
stop - start must be zero or have the same sign as step

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, 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, value: SIMD[dtype, 1], dims: List[Dim], location: Optional[_SourceLocation] = Optional(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]): 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(mut 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.

Implemented traits​

Methods​

__init__​

__getitem__​

debug_str​

__str__​

write_to​

verify​

layer​

current_layer​

nvop​

op​

constant​

quantize​

vector​

scalar​

range​

full​

output​

Implemented traits

Methods

`init`

`getitem`

`debug_str`

`str`

`write_to`

`verify`

`layer`

`current_layer`

`nvop`

`op`

`constant`

`quantize`

`vector`

`scalar`

`range`

`full`

`output`