Performance: einsum, ufunc tricks and avoiding copies

Express awkward tensor contractions readably, cut temporary allocations with in-place operations, and know when tuned C is worth it.

einsum notation

In einsum you write the index letters of the inputs, then an arrow and the letters you want to keep. Repeated letters are multiplied and summed; letters missing from the output are contracted away. Reading five such calls teaches more than any amount of axis juggling.

A = rng.random((100, 200))
B = rng.random((200, 50))
M = rng.random((4, 4))
T = rng.random((3, 4, 5))

np.einsum("ij,jk->ik", A, B)     # matrix product
np.einsum("ii->i", M)            # diagonal
np.einsum("ij->ji", A)           # transpose
np.einsum("ij,ij->", A, A)       # sum of all squares
np.einsum("ij,ij->i", A, A)      # one dot product per row
np.einsum("ij,kj->ik", A, B)     # row-wise similarity matrix
np.einsum("...ij->...ji", T)     # ellipsis: batch-aware transpose

Cutting temporary arrays

Every expression such as a + b * 2 allocates intermediate arrays. For large data the cost is memory traffic, not arithmetic, so eliminating temporaries is often the biggest single win.

a = np.ones(10_000_000)
b = np.ones(10_000_000)

a += b                      # in place, no new array
a *= 2.0

target = np.empty_like(a)
np.add(a, b, out=target)    # write into preallocated memory

np.einsum("ij,jk->ik", A, B, out=target2)   # out is a keyword argument

# conditional in-place update without a mask copy
np.copyto(target, a, where=target > 0.5)

# the layout rule of thumb: iterate over the last axis
m = rng.random((2000, 2000))
m.sum(axis=1)   # fast, contiguous
m.sum(axis=0)   # slower, strides down the buffer

Measure before reaching for a compiler

import timeit

timeit.timeit(lambda: A @ B, number=50) / 50     # seconds per call
%timeit A @ B                                    # in IPython and notebooks

# NumPy 2 has an array API and optional typing: profile first, then
# consider numexpr (fused element-wise expressions)
import numexpr as ne
ne.evaluate("2*a*a + 3*b", out=target)

# or Numba for genuine per-element logic that cannot be vectorised
from numba import njit
@njit(cache=True)
def rolling_sum(x, w):
    out = np.empty(x.size - w + 1)
    for i in range(out.size):
        out[i] = x[i:i + w].sum()
    return out

FAQ

Related

Views, copies and memory layout Linear algebra with numpy.linalg

Last refreshed 2026-09-18.