Mojo NumPy Integration
NumPy is Python's most-used numerical computing library. It provides multi-dimensional arrays, mathematical functions, linear algebra routines, and random number generation — all backed by optimized C and Fortran code. Mojo runs NumPy directly through its Python interop layer, letting you combine NumPy's rich ecosystem with Mojo's native performance for the computationally heavy parts.
The Collaboration Model
NumPy's strengths Mojo's strengths ───────────────────── ───────────────────────── Array creation and I/O SIMD vectorized loops Broadcasting rules Manual memory control Linear algebra (LAPACK) Zero-overhead parallelism Random number generation Compile-time specialization Plotting integration Direct hardware targeting Best strategy: Load data with NumPy → compute hot loops in Mojo → output with NumPy
Importing NumPy
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
print(np.__version__) # e.g. 1.26.4
Creating NumPy Arrays
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
# From a Python list
var a = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(a) # [1. 2. 3. 4. 5.]
print(a.shape) # (5,)
print(a.dtype) # float64
# 2D array (matrix)
var mat = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(mat)
print(mat.shape) # (3, 3)
# Zeros, ones, range
var zeros = np.zeros([4, 4])
var ones = np.ones([2, 3])
var rng = np.arange(0, 10, 2) # [0 2 4 6 8]
var space = np.linspace(0.0, 1.0, 5) # [0. 0.25 0.5 0.75 1.]
print(rng)
print(space)
NumPy Array Operations
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
var a = np.array([10.0, 20.0, 30.0, 40.0])
var b = np.array([ 1.0, 2.0, 3.0, 4.0])
# Element-wise arithmetic
print(a + b) # [11. 22. 33. 44.]
print(a * b) # [10. 40. 90. 160.]
print(a / b) # [10. 10. 10. 10.]
# Scalar broadcast
print(a * 2.0) # [20. 40. 60. 80.]
print(a + 5.0) # [15. 25. 35. 45.]
# Reductions
print(np.sum(a)) # 100.0
print(np.mean(a)) # 25.0
print(np.max(a)) # 40.0
print(np.min(a)) # 10.0
print(np.std(a)) # standard deviation
NumPy Linear Algebra
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
var A = np.array([[2.0, 1.0], [5.0, 3.0]])
var b = np.array([4.0, 7.0])
# Solve the linear system Ax = b
var x = np.linalg.solve(A, b)
print(x) # [5. -6.] — the solution
# Matrix multiply
var C = np.array([[1.0, 2.0], [3.0, 4.0]])
var D = np.array([[5.0, 6.0], [7.0, 8.0]])
print(np.matmul(C, D))
# [[19. 22.]
# [43. 50.]]
# Determinant and inverse
print(np.linalg.det(C)) # -2.0
print(np.linalg.inv(C))
Passing NumPy Arrays to Mojo Functions
The key performance pattern: use NumPy to set up data, then hand a raw buffer pointer to Mojo code for the heavy computation loop.
from python import Python, PythonObject
from memory import UnsafePointer
fn mojo_scale(data: UnsafePointer[Float64], n: Int, factor: Float64):
for i in range(n):
data[i] *= factor
fn main() raises:
var np = Python.import_module("numpy")
var ctypes = Python.import_module("ctypes")
# Create a NumPy array of float64
var arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float64)
print("Before:", arr) # [1. 2. 3. 4. 5.]
# Get a raw C pointer to the array's buffer
var n = int(arr.size)
var ptr_val = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
var raw = UnsafePointer[Float64].address_of(ptr_val[0])
# Run Mojo computation directly on NumPy's memory
mojo_scale(raw, n, 3.0)
print("After:", arr) # [3. 6. 9. 12. 15.] — NumPy array modified in-place
Memory sharing diagram:
NumPy array in Python
┌──────────────────────────────────┐
│ buffer: [ 1.0 | 2.0 | 3.0 ... ] │
│ ↑ │
│ data pointer │
└──────────────────────────────────┘
│
│ UnsafePointer points here
↓
Mojo mojo_scale() operates directly on the same memory.
No copy. No overhead.
NumPy Random Numbers
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
var rng = np.random.default_rng(seed=42)
# Uniform random floats in [0, 1)
var uniform = rng.random(size=5)
print(uniform)
# Normal distribution (mean=0, std=1)
var normal = rng.standard_normal(size=5)
print(normal)
# Random integers in [0, 100)
var ints = rng.integers(0, 100, size=5)
print(ints)
# Shuffle an array
var data = np.arange(10)
rng.shuffle(data)
print(data)
NumPy Indexing and Slicing
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
var m = np.arange(12).reshape([3, 4])
print(m)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]]
print(m[1, 2]) # 6 — row 1, col 2
print(m[0]) # [0 1 2 3] — first row
print(m[:, 1]) # [1 5 9] — second column
print(m[0:2, 1:3]) # [[1 2] [5 6]] — submatrix
print(m[m > 5]) # [6 7 8 9 10 11] — boolean mask
Saving and Loading NumPy Data
from python import Python
fn main() raises:
var np = Python.import_module("numpy")
# Save to binary .npy file
var weights = np.random.default_rng(0).standard_normal(size=[128, 64])
np.save("weights.npy", weights)
print("Saved:", weights.shape)
# Load back
var loaded = np.load("weights.npy")
print("Loaded:", loaded.shape) # (128, 64)
# Save multiple arrays to .npz archive
var biases = np.zeros(64)
np.savez("model.npz", weights=weights, biases=biases)
# Load from archive
var archive = np.load("model.npz")
print(archive["weights"].shape) # (128, 64)
print(archive["biases"].shape) # (64,)
Key Takeaways
Import NumPy with Python.import_module("numpy") inside any Mojo function. NumPy arrays support element-wise arithmetic, broadcasting, and reductions with one-line calls. Use np.linalg for solving linear systems, inverses, and determinants. Share memory between NumPy and Mojo by extracting the array's data pointer with ctypes and casting it to an UnsafePointer — zero copy, full Mojo speed on NumPy's buffer. Save arrays to .npy binary files for fast reloading and use .npz archives for multi-array checkpoints. The golden workflow: load and preprocess with NumPy, compute hot loops with Mojo, output and visualize with Python.