NumPy Reflexes That Carry Over (and Where Things Diverge)

The 90% that just works

If you have NumPy muscle memory, most of it transfers to MLX without you noticing. Indexing with integers, slicing with colons, broadcasting two differently-shaped arrays in an arithmetic op, reshape, the reductions (sum, mean, argmax) — all of these work the way you expect. You can write a surprising amount of MLX code by typing what you'd type in NumPy.

This lesson catalogs the carry-over so you stop hesitating, then names the two divergences that will trip you up if you don't see them coming.

What carries over (run the code, recognize everything)

Indexing returns a sub-array. Slicing returns a sub-array. Boolean masking returns a sub-array. Broadcasting works. Reductions work, with or without an axis. Reshape works. In-place mutation via item assignment works (this surprised me — early MLX docs implied otherwise; as of 0.31.x, a[i, j] = value is fine).

The two divergences worth memorizing

1. Lazy by default. When you write y = x * 2 + 1 in NumPy, the result is computed and stored. In MLX, the result is a graph node describing the computation; the actual numbers don't get computed until something forces evaluation. Anything that needs to read a concrete value triggers an "implicit eval" — printing the array, calling .item() on a scalar, converting to a Python list. The next lesson is dedicated to this; for now, the takeaway is that the array looks like NumPy but the timing of when computation happens is different.

2. Random API call signature is different. NumPy: np.random.randn(1024, 1024) — pass shape as separate positional args. MLX: mx.random.normal((1024, 1024)) — pass shape as a tuple. The same applies to uniform, randint, etc. Forget the tuple and you'll get a confusing error — not an obvious "shape must be a tuple" but a less helpful complaint about argument count.

That's actually it

Those are the two reflexes worth flagging. Everything else mostly carries over. If you find yourself typing what feels like NumPy code, you're probably writing valid MLX code; the times you trip will almost always be the lazy-eval timing or the random-API tuple thing. Mark this lesson and come back to it when one of those bites you.

Code

What carries over — recognize everything·python

import mlx.core as mx

a = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
print('a       :', a)
print('a[0]    :', a[0])
print('a[:, 1] :', a[:, 1])
print('a.sum() :', a.sum())
print('a.sum(axis=0):', a.sum(axis=0))
print('a.mean():', a.mean())
print('a.argmax(axis=1):', a.argmax(axis=1))

# Broadcasting (shape (3,) broadcasts against shape (2, 3))
b = mx.array([10.0, 20.0, 30.0])
print('a + b   :', a + b)

# Reshape
print('a.reshape(3,2):', a.reshape(3, 2))

# In-place via item assignment — works fine in mlx 0.31.x
a[0, 0] = 99.0
print('after a[0,0]=99 :', a)

Divergence 1: lazy by default (more in lesson 4)·python

import mlx.core as mx

x = mx.array([1.0, 2.0, 3.0])
y = x * 2 + 1                  # NumPy would compute now; MLX records a graph
# At this point y is a graph node, not concrete numbers.
print('y :', y)                 # ← print() triggers implicit eval; you see [3, 5, 7]

# To force computation explicitly, without a print:
big = mx.random.normal((1024, 1024))
big_squared = big @ big.T       # lazy — no work done yet
mx.eval(big_squared)            # NOW the kernel runs
print('big_squared.sum():', float(big_squared.sum()))

Divergence 2: random API takes a shape tuple·python

import mlx.core as mx

# NumPy:  np.random.randn(3, 4)
# MLX:    mx.random.normal((3, 4))    — shape is a single tuple argument

x = mx.random.normal((3, 4))
print('normal(3,4)  shape:', x.shape, 'dtype:', x.dtype)

u = mx.random.uniform(low=0.0, high=1.0, shape=(2, 3))
print('uniform(2,3) shape:', u.shape)

# Reproducibility: seed once at the top of your script
mx.random.seed(42)
print('seeded sample:', mx.random.normal((3,)))

Exercise

Take a NumPy snippet you've written before — anything you remember off the top of your head, ten lines or so — and translate it into MLX by typing the same thing with mx instead of np. When you hit a wall, the wall is almost certainly one of the two divergences (lazy or random-tuple). Time how long the translation takes and how many of your hits were the two divergences vs something else. Two sentences on what surprised you.