Eager execution
The max.experimental API abstracts the MAX
framework's graph foundation so you can write and run
models using PyTorch-like programming patterns. In other words, this API lets
you execute models eagerly, without thinking about graph compilation. The
fastest way to learn the API is to walk through the steps of building a small
forward pass.
Create a tensorβ
A Tensor is the
fundamental data structure for eager execution. This example creates a Tensor
with hard-coded data:
from max.driver import CPU
from max.experimental.tensor import Tensor
x = Tensor([1.0, -2.0, 3.0, -4.0, 5.0], device=CPU())
print(x)The expected output is:
Tensor([1 -2 3 -4 5], dtype=DType.float32, device=Device(type=cpu,id=0))Apply operations to tensorsβ
MAX lets you use Python operators (+, -, *, /, @) for arithmetic and
the functional API (F.relu(), F.softmax(), F.sqrt(), and so on) for more
complex computation.
Here's how to perform arithmetic operations using Python operators:
from max.experimental.tensor import Tensor
a = Tensor([1.0, 2.0, 3.0])
b = Tensor([4.0, 5.0, 6.0])
c = a + b # Addition
d = a * b # Element-wise multiplication
print(c)
print(d)Here are some
max.experimental.functional APIs
in use:
from max.experimental import functional as F
from max.driver import CPU
from max.experimental.tensor import Tensor
# Force CPU execution to avoid GPU compiler issues
x = Tensor([[1.0, 2.0], [3.0, 4.0]], device=CPU())
y = F.sqrt(x) # Element-wise square root
z = F.softmax(x, axis=-1) # Softmax along last axis
print(f"Input: {x}")
print(f"Square root: {y}")
print(f"Softmax: {z}")Inspect intermediate tensor valuesβ
One of the biggest advantages of the eager API is that you can inspect intermediate values at any point in your code. Use this capability to verify that intermediate computations match your expectations as you build up a forward pass.
from max.experimental import functional as F
from max.driver import CPU
from max.dtype import DType
from max.experimental.tensor import Tensor
def debug_forward_pass(x: Tensor) -> Tensor:
"""Forward pass with intermediate inspection."""
print(f"Input: {x}")
z = x * 2
print(f"After multiply: {z}")
h = F.relu(z)
print(f"After ReLU: {h}")
return h
x = Tensor([-1.0, 0.0, 1.0, 2.0], dtype=DType.float32, device=CPU())
result = debug_forward_pass(x)The expected output is:
Input: Tensor([-1 0 1 2], dtype=DType.float32, device=Device(type=cpu,id=0))
After multiply: Tensor([-2 0 2 4], dtype=DType.float32, device=Device(type=cpu,id=0))
After ReLU: Tensor([0 0 2 4], dtype=DType.float32, device=Device(type=cpu,id=0))Run a forward passβ
A forward pass strings together everything you've seen so far into the shape of a single neural network layer. Here's a realistic example with random weights, a matmul, an activation, and a reduction.
from max.experimental import functional as F
from max.experimental import random
from max.driver import CPU
from max.dtype import DType
from max.experimental.tensor import Tensor
# Create input data
x = Tensor([[1.0, 2.0], [3.0, 4.0]], dtype=DType.float32, device=CPU())
# Create random weights
w = random.gaussian(
[2, 2], mean=0.0, std=0.1, dtype=DType.float32, device=CPU()
)
# Run forward pass
z = x @ w # Matrix multiply
h = F.relu(z) # Activation
out = h.mean() # Reduce to scalar
# Inspect results
print(f"Input shape: {x.shape}")
print(f"After matmul: {z.shape}")
print(f"Output: {out}")Shape and type validation happens as you write operations. If the matrix
dimensions didn't align for x @ w, you'd get an error at that exact line
showing the shape mismatch. Try modifying the code to see how validation
errors appear at the line that caused them.
Next stepsβ
Now that you know how to write a small forward pass, explore more deeply about the building blocks you'll use day-to-day:
- Tensor fundamentals: Learn what tensors are, how to create them, and how to inspect their attributes.
- Basic operations: Apply arithmetic operators, reductions, and the functional API to tensors.
- Data types: Specify how each element in a tensor is represented in memory.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!