Get started with MAX Graph in Python

Modular Team

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 example-project && cd example-project

   mkdir example-project && cd example-project
2. Create and activate a virtual environment:

   python3 -m venv .venv/example-project \
     && source .venv/example-project/bin/activate

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

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

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

   Stable

   pip install modular \
     --extra-index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://modular.gateway.scarf.sh/simple/

   pip install modular \
     --extra-index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://modular.gateway.scarf.sh/simple/

uv

1. If you don't have it, 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 example-project && cd example-project

   uv init example-project && cd example-project
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 \
     --extra-index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://dl.modular.com/public/nightly/python/simple/ \
     --index-strategy unsafe-best-match

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

   Stable

   uv pip install modular \
     --extra-index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://modular.gateway.scarf.sh/simple/ \
     --index-strategy unsafe-best-match

   uv pip install modular \
     --extra-index-url https://download.pytorch.org/whl/cpu \
     --extra-index-url https://modular.gateway.scarf.sh/simple/ \
     --index-strategy unsafe-best-match

conda

1. If you don't have it, install conda. A common choice is with brew:

   brew install miniconda

   brew install miniconda
2. Initialize conda for shell interaction:

   conda init

   conda init

   If you're on a Mac, instead use:

   conda init zsh

   conda init zsh

   Then restart your terminal for the changes to take effect.

3. Create a project:

   conda create -n example-project

   conda create -n example-project
4. Start the virtual environment:

   conda activate example-project

   conda activate example-project
5. Install the modular conda package:
   Nightly

   conda install -c conda-forge -c https://conda.modular.com/max-nightly/ modular

   conda install -c conda-forge -c https://conda.modular.com/max-nightly/ modular

   Stable

   conda install -c conda-forge -c https://conda.modular.com/max/ modular

   conda install -c conda-forge -c https://conda.modular.com/max/ modular

pixi

1. If you don't have it, install pixi:

   curl -fsSL https://pixi.sh/install.sh | sh

   curl -fsSL https://pixi.sh/install.sh | sh

   Then restart your terminal for the changes to take effect.

2. Create a project:

   pixi init example-project \
     -c https://conda.modular.com/max-nightly/ -c conda-forge \
     && cd example-project

   pixi init example-project \
     -c https://conda.modular.com/max-nightly/ -c conda-forge \
     && cd example-project
3. Install the modular conda package:
   Nightly

   pixi add modular

   pixi add modular

   Stable

   pixi add "modular==25.3"

   pixi add "modular==25.3"
4. Start the virtual environment:

   pixi shell

   pixi 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

pixi

Change folders to your working directory:

cd src/quickstart

cd src/quickstart

You can check your MAX version like this:

pixi run max --version

pixi run max --version

You can check your Python version like this:

pixi run python --version

pixi 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

pixi

pixi run python addition.py

pixi 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