Get started with MAX Graph in Python

Patrick Rachford

Technical Writer

max-graph

python

numpy

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:

1. Build a simple graph that adds two numbers
2. Create an inference session to load and compile the graph
3. 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.

Set up your environment

Create a Python project to install our APIs and CLI tools.

pip

1. Create a project folder:

   mkdir modular && cd modular

   mkdir modular && cd modular
2. Create and activate a virtual environment:

   python3 -m venv .venv/modular \
     && source .venv/modular/bin/activate

   python3 -m venv .venv/modular \
     && source .venv/modular/bin/activate
3. Install the modular Python package:
   Nightly

   pip install modular \
     --index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://dl.modular.com/public/nightly/python/simple/

   pip install modular \
     --index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://dl.modular.com/public/nightly/python/simple/

   Stable

   pip install modular \
     --index-url https://download.pytorch.org/whl/cpu

   pip install modular \
     --index-url https://download.pytorch.org/whl/cpu

uv

1. Install uv:

   curl -LsSf https://astral.sh/uv/install.sh | sh

   curl -LsSf https://astral.sh/uv/install.sh | sh

   Then restart your terminal to make uv accessible.

2. Create a project:

   uv init modular && cd modular

   uv init modular && cd modular
3. Create and start a virtual environment:

   uv venv && source .venv/bin/activate

   uv venv && source .venv/bin/activate
4. Install the modular Python package:
   Nightly

   uv pip install modular \
     --index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://dl.modular.com/public/nightly/python/simple/

   uv pip install modular \
     --index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://dl.modular.com/public/nightly/python/simple/

   Stable

   uv pip install modular \
     --index-url https://download.pytorch.org/whl/cpu

   uv pip install modular \
     --index-url https://download.pytorch.org/whl/cpu

magic

1. Install magic:

   curl -ssL https://magic.modular.com/ | bash

   curl -ssL https://magic.modular.com/ | bash

   Then run the source command that's printed in your terminal.

2. Create a project:

   magic init modular --format pyproject && cd modular

   magic init modular --format pyproject && cd modular
3. Install the max-pipelines conda package:
   Nightly

   magic add max-pipelines

   magic add max-pipelines

   Stable

   magic add "max-pipelines==25.3"

   magic add "max-pipelines==25.3"
4. Start the virtual environment:

   magic shell

   magic shell

Then, create a working directory.

pip

Create a folder called max_ops:

mkdir max_ops
cd max_ops

mkdir max_ops
cd max_ops

You can check your MAX version like this:

pip show modular

pip show modular

You can check your Python version like this:

python --version

python --version

uv

Create a folder called max_ops:

mkdir max_ops
cd max_ops

mkdir max_ops
cd max_ops

You can check your MAX version like this:

uv pip show modular

uv pip show modular

You can check your Python version like this:

python --version

python --version

magic

Change folders to your working directory:

cd src/quickstart

cd src/quickstart

You can check your MAX version like this:

magic run max --version

magic run max --version

You can check your Python version like this:

magic run python --version

magic run python --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 your working directory and add the following libraries:

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

from typing import Any
import numpy as np
from max import engine
from max.dtype import DType
from max.graph import DeviceRef, 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,), device=DeviceRef.CPU()
    )
    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,), device=DeviceRef.CPU()
    )
    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 execution, rather than containing the actual 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(a, b)[0]
    print("result:", ret)
    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(a, b)[0]
    print("result:", ret)
    return ret

Starting in 24.6.0, the model.execute() command no longer accepts keyword arguments. In a future release we will restore this functionality with support for GPUs. For compatibility with existing code that uses keyword arguments, you can use the execute_legacy() function.

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, run the Python file from the command line:

pip

python addition.py

python addition.py

uv

python addition.py

python addition.py

magic

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:

• MAX Graph API example
• MAX Graph implementation of LLama3

Next steps

Did this tutorial work for you?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It's missing a step; I don't know what to do

I don't understand the purpose

Other

Submit