Get started with MAX Graph in Python

AI Technical Writer

Updated: Nov 2, 2024 |

14 min read

python

graph

MAX Graph is a high-performance computation framework that lets you build and execute efficient machine learning models. It provides a flexible way to define computational workflows as graphs, where each node represents an operation (like matrix multiplication or addition) and edges represent the flow of data. By using MAX Graph, you can create optimized machine learning models that run faster and more efficiently on modern hardware.

In this tutorial, you'll build a graph using the MAX Graph API in Python with an ops function.

To do this, you will complete the following steps:

Build a simple graph that adds two numbers
Create an inference session to load and compile the graph
Execute the graph with input data

By the end of this tutorial, you'll have an understanding of how to construct basic computational graphs, set up inference sessions, and run computations using the MAX Graph API.

Create a virtual environment

Use the Magic CLI to create the environment and install the required packages.

🪄

If you don't have Magic, click here.

Install it with one command.

Create a project with Python and change into the max_ops directory:

magic init max_ops --format pyproject
cd max_ops
magic init max_ops --format pyproject
cd max_ops

Then add your project dependency packages:

magic add "max~=24.5" "numpy<2.0"
magic add "max~=24.5" "numpy<2.0"

You can check your Python version like this:

magic run python3 --version
magic run python3 --version

If you have any questions along the way, ask them on our Discord channel.

1. Build the graph

Now with our environment and packages setup, lets create the graph. This graph will define a computational workflow that adds two tensors together.

Let's start by creating a new file called addition.py inside of the src/max_ops folder and add the following libraries:

import numpy as np
from max import engine
from max.dtype import DType
from max.graph import Graph, TensorType, ops
import numpy as np
from max import engine
from max.dtype import DType
from max.graph import Graph, TensorType, ops

To create a computational graph, use the Graph() class from the MAX Graph API. When initializing, specify a name for the graph and define the types of inputs it will accept.

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    input_type = TensorType(dtype=DType.float32, shape=(1,))
    with Graph(
        "simple_add_graph", input_types=(input_type, input_type)
    ) as graph:
        lhs, rhs = graph.inputs
        out = ops.add(lhs, rhs)
        graph.output(out)
def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    input_type = TensorType(dtype=DType.float32, shape=(1,))
    with Graph(
        "simple_add_graph", input_types=(input_type, input_type)
    ) as graph:
        lhs, rhs = graph.inputs
        out = ops.add(lhs, rhs)
        graph.output(out)

Inside the context manager, access the graph's inputs using the inputs property. This returns a symbolic tensor representing the input arguments.

The symbolic tensor is a placeholder that represents the shape and type of data that will flow through the graph during the exectuion, rather than containing the acutal numeric values like in eager execution.

Then use the add() function from the ops package to add the two input tensors. This creates a new symbolic tensor representing the sum.

Finally, set the output of the graph using the output() method. This specifies which tensors should be returned when the graph is executed.

Now, add a print() function to the graph to see what's created.

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
        # ...
        print("final graph:", graph)
def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
        # ...
        print("final graph:", graph)

The output will show us the structure of our graph, including the input it expects and the operations it will perform. This helps us understand how our graph will process data when we use it.

Next, let's load the graph into an inference session.

2. Create an inference session

Now that our graph is constructed, let's set up an environment where it can operate. This involves creating an inference session and loading our graph into it.

Create an InferenceSession() instance that loads and runs the graph inside the add_tensors() function.

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    # ... 
    # 2. Create an inference session
    session = engine.InferenceSession()
    model = session.load(graph)
def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    # ... 
    # 2. Create an inference session
    session = engine.InferenceSession()
    model = session.load(graph)

This step transforms our abstract graph into a computational model that's ready for execution.

To ensure our model is set up correctly, let's examine its input requirements.

Print the graph's input metadata by using the input_metadata property.

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    # ... 
    # 2. Create an inference session
    session = engine.InferenceSession()
    model = session.load(graph)
    for tensor in model.input_metadata:
        print(
            f"name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}"
        )
def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # 1. Build the graph
    # ... 
    # 2. Create an inference session
    session = engine.InferenceSession()
    model = session.load(graph)
    for tensor in model.input_metadata:
        print(
            f"name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}"
        )

This will output the exact specifications of the input our model expects, helping us prepare appropriate data for processing.

Next, let's execute the graph.

3. Execute the graph

To give the model something to add, create two inputs of a shape and a data type that match our graph's input requirements. Then pass the inputs to the execute() function:

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # ... 
    # 2. Create an inference session
    # ... 
    # 3. Execute the graph
    ret = model.execute(input0=a, input1=b)
    print("result:", ret["output0"])
    return ret
def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, any]:
    # ... 
    # 2. Create an inference session
    # ... 
    # 3. Execute the graph
    ret = model.execute(input0=a, input1=b)
    print("result:", ret["output0"])
    return ret

4. Run the example

Now that we've built our graph, created an inference session, and defined how to execute the graph, let's put it all together and run our complete example.

At the end of your addition.py file, add the following code:

if __name__ == "__main__":
    input0 = np.array([1.0], dtype=np.float32)
    input1 = np.array([1.0], dtype=np.float32)
    add_tensors(input0, input1)
if __name__ == "__main__":
    input0 = np.array([1.0], dtype=np.float32)
    input1 = np.array([1.0], dtype=np.float32)
    add_tensors(input0, input1)

This passes your arguments input0 and input1 to the add_tensors() function.

Then using magic run command on the Python file from the command line:

magic run python addition.py

magic run python addition.py

You've successfully created your first graph using the MAX Graph API in Python. Let's examine what was printed to the terminal:

final graph: mo.graph @simple_add_graph(%arg0: !mo.tensor<[1], f32>, %arg1: !mo.tensor<[1], f32>) -> !mo.tensor<[1], f32> attributes {argument_names = ["input0", "input1"], result_names = ["output0"]} {
  %0 = rmo.add(%arg0, %arg1) : (!mo.tensor<[1], f32>, !mo.tensor<[1], f32>) -> !mo.tensor<[1], f32>
  mo.output %0 : !mo.tensor<[1], f32>
}
final graph: mo.graph @simple_add_graph(%arg0: !mo.tensor<[1], f32>, %arg1: !mo.tensor<[1], f32>) -> !mo.tensor<[1], f32> attributes {argument_names = ["input0", "input1"], result_names = ["output0"]} {
  %0 = rmo.add(%arg0, %arg1) : (!mo.tensor<[1], f32>, !mo.tensor<[1], f32>) -> !mo.tensor<[1], f32>
  mo.output %0 : !mo.tensor<[1], f32>
}

Two input tensors (%arg0, %arg1) of shape [1] and float32 type
The addition operation connecting them
One output tensor of matching shape/type

The metadata lines confirm both input tensors match the required specifications.

name: input0, shape: [1], dtype: DType.float32
name: input1, shape: [1], dtype: DType.float32
name: input0, shape: [1], dtype: DType.float32
name: input1, shape: [1], dtype: DType.float32

The result shows the addition worked correctly:

[1.0] + [1.0] = [2.0]

result: [2.]

result: [2.]

Now that you've built your first MAX Graph that performs addition, you can explore more complex examples:

Create a virtual environment​

1. Build the graph​

2. Create an inference session​

3. Execute the graph​

4. Run the example​

Create a virtual environment

1. Build the graph

2. Create an inference session

3. Execute the graph

4. Run the example