Lazy Evaluation — The One Mental Model Shift That Matters

The shift, in one paragraph

In NumPy and PyTorch's eager mode, every line of array code does the work right then. y = x * 2 + 1 computes and stores the result; if you never use y, you wasted the work. In MLX, that same line builds a tiny graph describing the computation. The work happens later — when you (explicitly or implicitly) ask for the value.

That's the entire shift. Everything else about lazy evaluation in MLX is consequences of this one decision.

What triggers actual computation

Explicit: mx.eval(...) on one or more arrays — the canonical way to materialize.

Implicit: anything that needs to read a concrete value. print(array) triggers eval (so you see numbers, not a graph). .item() on a scalar triggers eval. Converting to a NumPy array or Python list triggers eval. Indexing into the array also triggers eval, because indexing needs to know the values.

This is why MLX feels like NumPy when you're poking at small examples in a REPL — every print is silently forcing the work, so the laziness is invisible. It only matters when you string many operations together and want them to fuse into fewer GPU dispatches.

Why MLX bothers with this

Kernel fusion — when MLX sees a graph of operations, it can sometimes fuse them into a single Metal kernel. mx.compile (lesson 7) leans on this. Eager mode can't.
Function transforms — mx.grad, mx.vmap work cleanly because they get the whole function as a graph, not a sequence of side-effecting calls. JAX makes the same trade.
Memory efficiency — intermediates that are never read can be skipped entirely.

The debugging implication

Because eval is implicit when you print, your "step-through" intuition from NumPy still works at the REPL — you'll see numbers when you ask. The trap is that production code without prints can keep things lazy across a much longer sequence than you expect, which means errors don't surface until mx.eval at the end of the chain. Get used to two debugging idioms — sprinkle mx.eval on suspect intermediates while you triage, and sprinkle print for the same effect when you're at the REPL.

Code

Eager-looking, lazy-acting·python

import mlx.core as mx

x = mx.array([1.0, 2.0, 3.0])
y = x * 2 + 1                  # builds a graph; no kernel has run
# At this point, y is a graph node, not concrete numbers.

# Implicit eval — print materializes
print('y :', y)                 # → array([3, 5, 7], dtype=float32)

Explicit eval — for big sequences without prints·python

import mlx.core as mx

a = mx.random.normal((1024, 1024))
b = mx.random.normal((1024, 1024))

# All three lines below build a graph; nothing executes yet.
c = a @ b
d = c + a
e = mx.tanh(d)

# Force the whole graph to execute, in one go (better fusion opportunity).
mx.eval(e)

print('e mean:', float(e.mean()))   # implicit eval, but e is already materialized

Multiple outputs in one mx.eval call·python

import mlx.core as mx

x = mx.array([1.0, 2.0, 3.0])
sq = x ** 2
sm = sq.sum()
mn = sq.mean()

# Materialize all three at once — MLX can schedule them efficiently together.
mx.eval(sq, sm, mn)
print('sq:', sq, 'sum:', float(sm), 'mean:', float(mn))

Exercise

Construct a chain of five operations on a (512, 512) random array (e.g., matmul, add, tanh, sum, mean) and time two versions: (a) one big mx.eval at the end, (b) mx.eval after every line. Use time.perf_counter(), run each three times, take the median. Are they meaningfully different on your machine? They should be — the first version gets fusion benefits the second one doesn't.