
Get started with MAX Graph in Python
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.
Set up your environment
Create a Python project to install our APIs and CLI tools.
- pip
- uv
- conda
- pixi
- Create a project folder:
mkdir example-project && cd example-project
mkdir example-project && cd example-project
- Create and activate a virtual environment:
python3 -m venv .venv/example-project \
&& source .venv/example-project/bin/activatepython3 -m venv .venv/example-project \
&& source .venv/example-project/bin/activate - Install the
modular
Python package:- Nightly
- Stable
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/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/
- 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. - Create a project:
uv init example-project && cd example-project
uv init example-project && cd example-project
- Create and start a virtual environment:
uv venv && source .venv/bin/activate
uv venv && source .venv/bin/activate
- Install the
modular
Python package:- Nightly
- Stable
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-matchuv 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-matchuv 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-matchuv 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
- If you don't have it, install conda. A common choice is with
brew
:brew install miniconda
brew install miniconda
- 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.
- Create a project:
conda create -n example-project
conda create -n example-project
- Start the virtual environment:
conda activate example-project
conda activate example-project
- Install the
modular
conda package:- Nightly
- Stable
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
conda install -c conda-forge -c https://conda.modular.com/max/ modular
conda install -c conda-forge -c https://conda.modular.com/max/ modular
- 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.
- Create a project:
pixi init example-project \
-c https://conda.modular.com/max-nightly/ -c conda-forge \
&& cd example-projectpixi init example-project \
-c https://conda.modular.com/max-nightly/ -c conda-forge \
&& cd example-project - Install the
modular
conda package:- Nightly
- Stable
pixi add modular
pixi add modular
pixi add "modular==25.3"
pixi add "modular==25.3"
- Start the virtual environment:
pixi shell
pixi shell
Then, create a working directory.
- pip
- uv
- pixi
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
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
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
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
- uv
- pixi
python addition.py
python addition.py
python addition.py
python addition.py
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:
result: [2.]
result: [2.]
Now that you've built your first MAX Graph that performs addition, you can explore more complex examples:
Next steps
Did this tutorial work for you?
Thank you! We'll create more content like this.
Thank you for helping us improve!