Skip to main content

graph

Core graph primitives.

Graphs are callable routines in MAX Engine, similar to functions in Mojo. Like functions, graphs have a name and signature. Unlike functions, which follow an imperative programming model, Graphs follow a dataflow programming model, using lazily-executed, parallel operations instead of sequential instructions. Graphs aren't called directly from Mojo, but are instead compiled and executed by MAX Engine, for example using the MAX Engine API.

Graph

Represents a single MAX graph.

Implemented traits:

AnyType, CollectionElement, Copyable, Movable, Stringable

Methods:

__init__

__init__(inout self: Self, name: String, in_types: TypeTuple, out_types: TypeTuple)

Constructs a new Graph.

The constructed Graph will not be valid unless it has no outputs; a graph with outputs will need a graph.output call to tell it what to return. The graph's validity can be checked by calling graph.module().verify().

Args:

  • name (String): A name for the graph.
  • in_types (TypeTuple): The Graph's input types.
  • out_types (TypeTuple): The Graph's output types.

__getitem__

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

Returns the nth argument of this Graph.

The Symbol that the argument is returned as can be used as input to other Graph Nodes.

Raises: If n is not a valid argument number.

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.

__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.

module

module(self: Self) -> Module

Returns the Graph's parent Modules.

Returns:

The Module that holds this Graph.

nvop

nvop(self: Self, name: String, inputs: SymbolTuple, out_types: TypeTuple, attrs: AttrMap) -> SymbolTuple

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 module.

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 (SymbolTuple): The list of symbolic operands.
  • out_types (TypeTuple): The list of output types.
  • attrs (AttrMap): Any attributes that the operation might require.

Returns:

The symbolic outputs of the newly-added Node.

op

op(self: Self, name: String, out_type: AnyMOType, attrs: AttrMap) -> 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 (AnyMOType): The output types.
  • attrs (AttrMap): Any attributes that the operation might require.

Returns:

The symbolic output of the newly-added Node.

op(self: Self, name: String, inputs: SymbolTuple, out_type: AnyMOType, attrs: AttrMap) -> 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 (SymbolTuple): The list of symbolic operands.
  • out_type (AnyMOType): The output types.
  • attrs (AttrMap): 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 MOTensor 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 MOTensor 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) -> Symbol

Adds a Node representing a mo.constant operation.

The value of this constant will have the type scalar MOTensor (0-D shape), when rank is 0, or a higher-rank MOTensor 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: ElementType) -> Symbol

Adds a Node representing a mo.constant operation.

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

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

Args:

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

Returns:

The symbolic output of this Node.

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

Adds a Node representing a mo.constant operation.

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

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

Args:

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

Returns:

The symbolic output of this Node.

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: SymbolTuple) -> 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 (SymbolTuple): 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(self: Self, outs: SymbolTuple)

Adds a Node representing a mo.output operation.

This is a special Node that all Graphs must have. The inputs must match the Graphs signature (specifically, its return values).

Args:

  • outs (SymbolTuple): The Graphs return values.