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.
Create a virtual environment
Use the Magic CLI to create the environment and install the required packages.
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_legacy(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_legacy(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:
result: [2.]
result: [2.]
Now that you've built your first MAX Graph that performs addition, you can explore more complex examples:
Did this tutorial work for you?
Thank you! We'll create more content like this.
Thank you for helping us improve!