Skip to main content

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?