Vectorization: Measure Speed, Shapes and Memory

The story — measure the loop you actually have

DL011 already batches its affine maps. Replacing that code with the same X@W cannot magically make it100× faster. The useful comparison is a genuinely scalar/per-example implementation against an equivalent bulk array implementation, with shape, dtype, thread count and numerical tolerance held fixed. This lesson measures dot products, MLP forward/backward, memory-layout effects and a vectorized version of DL012's four-point network.

NumPy moves loops into compiled numeric code; optimized BLAS can reuse tiled data, exploit SIMD and parallelize suitable matrix products. There is no universal100–1000× ratio, no requirement that every array be contiguous, and no fixed fraction of hardware peak for all shapes. Looping over layers, minibatches or bounded chunks is often exactly right.

Prerequisites and outcomes

Bring matrix shapes, broadcasting and the NumPy MLP. You will translate coordinate sums to einsum, compare scalar and batch programs, count FLOPs with a stated convention, inspect temporary memory, and identify sequential dependencies. Optional pre-reading: Andrew Ng, more examples, NumPy2.2 einsum, and ajcr's guide.

1 · Interpreter overhead, layout and arithmetic intensity

A Python loop over NumPy scalars pays per-iteration dispatch and scalar handling. Bulk ufuncs reduce that overhead, but separate expressions can allocate and traverse multiple temporary arrays. A strided array is still valid; optimized kernels may handle strides, pack data or copy internally. BLAS selection and threading depend on the installed build, not the spelling “NumPy.”

The freight-train analogy is about amortized dispatch and data reuse, not a million arithmetic operations executing in one clock. SIMD width is architecture/dtype dependent. A tiny vector can be dominated by function overhead while a large reduction becomes memory-bound.

2 · A concrete timing experiment

Use float64, a fixed generator, warmups, perf_counter, multiple measurements and a numerical assertion. The result is a measurement of this process, not a laptop promise. Set BLAS/OMP thread limits before starting Python for a controlled CPU comparison.

import numpy as np
import time
import statistics
import random
 
def dot_python(a,b):
    s=0.
    for i in range(len(a)):s+=a[i]*b[i]
    return s
 
def bench(fn,repeats=5):
    fn()
    times=[]
    for _ in range(repeats):
        t=time.perf_counter();fn();times.append(time.perf_counter()-t)
    return {'median_seconds':statistics.median(times),'samples_seconds':times}
 
rng=np.random.default_rng(13)
a,b=rng.normal(size=10000),rng.normal(size=10000)
ref=dot_python(a,b)
assert np.isclose(ref,a@b,rtol=1e-12,atol=1e-12)
assert np.isclose(ref,(a*b).sum(),rtol=1e-12,atol=1e-12)
print('dot timings', {name:bench(fn) for name,fn in [
    ('python',lambda:dot_python(a,b)),('temporary',lambda:(a*b).sum()),('dot',lambda:a@b)]})

Repeat for sizes10,10,000 and100,000 within your memory/time budget. The ratio can change because fixed overhead, cache residency and memory traffic change. Reordered floating-point sums need not be bit-identical; cancellation-heavy data deserve absolute as well as relative error checks. This is why “same math, exactly the same bits” is too strong.

3 · Einsum — output labels decide what is reduced

In explicit notation, labels absent from the output are summed. A repeated label is not automatically reduced: i,i->i retains i and multiplies elementwise. A repeated label within one operand can select a diagonal. Use ... for unspecified broadcast dimensions; inspect the resulting shape.

3.1 Six examples, executable rather than undefined names

A=rng.normal(size=(3,4));B=rng.normal(size=(4,2))
BA=rng.normal(size=(2,3,4));BB=rng.normal(size=(2,4,5))
assert np.allclose(np.einsum('i,i->',a,b),a@b)
assert np.allclose(np.einsum('i,i->i',a,b),a*b)
assert np.allclose(np.einsum('i,j->ij',a[:3],b[:4]),a[:3,None]*b[None,:4])
assert np.allclose(np.einsum('mk,kn->mn',A,B),A@B)
assert np.allclose(np.einsum('bmk,bkn->bmn',BA,BB),BA@BB)
Q=rng.normal(size=(2,3,4));K=rng.normal(size=(2,5,4));V=rng.normal(size=(2,5,6))
scores=np.einsum('bqd,bkd->bqk',Q,K)/np.sqrt(4)
shift=scores-scores.max(-1,keepdims=True)
weights=np.exp(shift);weights/=weights.sum(-1,keepdims=True)
out=np.einsum('bqk,bkv->bqv',weights,V)
assert out.shape==(2,3,6)
assert np.allclose(scores,Q@K.swapaxes(-1,-2)/2)
assert np.allclose(out,weights@V)
# dot is NOT batched matmul on these ranks.
assert np.dot(BA,BB).shape==(2,3,2,5)
assert (BA@BB).shape==(2,3,5)

Attention's value width can differ from key/query width, hence v in the final output. Softmax and nonlinearities are not expressible merely as a multilinear einsum contraction. The former “every transformer op is einsum” claim conflated contractions with the whole computation.

3.2 Contraction order and dispatch

np.einsum defaults to optimize=False; even a two-operand einsum need not use the same kernel or achieve the same speed as @. For several operands inspect np.einsum_path and then pass the chosen path or optimize='greedy'. optimal path search can itself be expensive and does not guarantee best wall time on every machine. Record path planning separately if reusing a contraction many times. PyTorch's optional opt_einsum integration is a different library contract, not an automatic NumPy cache.

4 · Vectorizing the MLP — before and after

The loop reference derives dimensions from weights rather than hardcoding128 and10. Both versions perform the same affine-ReLU-affine map. The vectorized version contains several ufunc/matmul calls; “one NumPy expression” is not “one kernel.”

def forward_naive(X,W1,b1,W2,b2):
    H=np.zeros((len(X),W1.shape[1]),dtype=X.dtype)
    for i in range(len(X)):
        for j in range(W1.shape[1]):H[i,j]=max(0,X[i]@W1[:,j]+b1[j])
    Z=np.zeros((len(X),W2.shape[1]),dtype=X.dtype)
    for i in range(len(X)):
        for j in range(W2.shape[1]):Z[i,j]=H[i]@W2[:,j]+b2[j]
    return Z
 
def forward_fast(X,W1,b1,W2,b2):
    return np.maximum(0,X@W1+b1)@W2+b2
 
def forward_einsum(X,W1,b1,W2,b2):
    H=np.maximum(0,np.einsum('nd,dh->nh',X,W1,optimize=True)+b1)
    return np.einsum('nh,hc->nc',H,W2,optimize=True)+b2
 
X=rng.normal(size=(16,32))
p=(rng.normal(size=(32,12)),rng.normal(size=12),rng.normal(size=(12,5)),rng.normal(size=5))
reference=forward_naive(X,*p)
assert np.allclose(reference,forward_fast(X,*p))
assert np.allclose(reference,forward_einsum(X,*p))
print('MLP timings',{f.__name__:bench(lambda:f(X,*p)) for f in (forward_naive,forward_fast,forward_einsum)})

4.1 The naive version

A loop over16 examples and12 hidden units makes192 Python hidden-unit iterations; output work adds80. It already uses a compiled dot for each unit, so it is not the same baseline as a three-level scalar multiplication loop. Name the baseline before reporting a speedup.

4.2 The vectorized version

Rows share a matrix multiply, improving reuse and amortizing dispatch. This changes rounding order but not the intended real-valued function. Assert outputs before timing; a faster wrong-axis result is not an optimization.

4.3 Timing interpretation

Warmup, thread count, process contention and allocation are part of the result. The audit uses small CPU tensors; it does not extrapolate a measured forward ratio to a four-hour versus nine-second MNIST epoch. Data loading, backward and optimizer work must be timed separately or end-to-end.

5 · Batching and sequential work

A larger batch can amortize launch overhead and expose parallel matrix work. An accelerator does not assign one complete example to a single SIMD lane, and neither batch64 nor4096 is a universal minimum/maximum. Tensor layouts differ; batch-first is a convention, not a law of every API.

At inference, request batching trades throughput against queueing latency and memory. Autoregressive sampling depends on prior sampled tokens; teacher-forced training can process known token positions in parallel with a causal mask. Generic nonlinear RNN recurrence has a sequential dependency, while some structured recurrences admit parallel scans. Variable-length beam/rejection work can batch active paths, at the cost of masking or compaction. Look for independent axes and avoid claiming that every temporal loop can—or can never—be parallelized.

6 · FLOP counting — state the convention

For (M,K)@(K,N), the conventional multiply-plus-add estimate is 2MKN; the exact scalar count for ordinary dot accumulation is MN(2K−1). Counting a fused multiply-add as two FLOPs is a convention, not two machine instructions.

For64×784 times784×128 the estimate is 12,845,056 FLOPs. For64×128 times128×10 it is 163,840, a ratio78.4. Dividing by an assumed100GFLOP/s gives a theoretical compute-only lower estimate of.12845ms for the first operation; it is not a measurement or a statement about integer precision.

6.1 Transformer forward cost

For a dense conventional block with width D and FFN width4D, attention projections cost about8BTD² FLOPs and the FFN16BTD². Dense score and weighted-value products add about4BT²D. Thus a rough sequence total is L*(24BTD²+4BT²D), ignoring norms, activations and vocabulary projection. Per token divides by BT. The original 12LD²T per token mixed MACs/FLOPs and sequence/per-token units. The common 2P per token estimate counts dense weight matmuls and omits important context-dependent attention, embeddings and sparsity effects; “7B→14GFLOPs” is that approximation, not all inference cost.

7 · Memory traffic, layout and loops worth keeping

Float32 elementwise addition performs n additions and minimally reads2n+writes n values: ideal intensity1/12 FLOP per byte, ignoring write allocation/cache effects. An n×n matmul does about2n³ FLOPs with an ideal3n²-value traffic floor, giving n/6 FLOP per byte for float32. Finite cache and packing increase actual traffic; small or skinny matmuls can still be memory/overhead-bound. The roofline estimate is min(peak_FLOPs,bandwidth*intensity).

(a+b)*c normally creates one intermediate and two arithmetic passes, not three. Adding a final sum adds a reduction pass. Fusion can reduce traffic, but its speedup is not guaranteed proportional to removed passes. A transpose is usually a metadata view; whether downstream traversal costs more depends on strides and kernel choices. In particular a.sum(0) and a.T.sum(1) describe the same underlying traversal opportunity, so they do not prove a layout penalty.

C=rng.normal(size=(256,384));Farr=np.asfortranarray(C)
assert np.allclose(C.sum(1),Farr.sum(1))
assert np.allclose(C.sum(0),C.T.sum(1))
print('layout',C.strides,Farr.strides,{name:bench(lambda:arr.sum(1)) for name,arr in [('C',C),('F',Farr)]})
# Broadcasting can allocate a large all-pairs temporary; compute in chunks.
def squared_distances_blocked(A,B,block=8):
    out=np.empty((len(A),len(B)))
    for s in range(0,len(A),block):
        diff=A[s:s+block,None,:]-B[None,:,:]
        out[s:s+block]=np.einsum('ijk,ijk->ij',diff,diff)
    return out
AA=rng.normal(size=(19,7));BB=rng.normal(size=(13,7))
assert np.allclose(squared_distances_blocked(AA,BB),((AA[:,None]-BB[None])**2).sum(-1))
assert not np.shares_memory(AA[[0,1]],AA)
assert np.shares_memory(AA[:2],AA)

Broadcasting can be readable but materialize O(NMD) memory. A contraction/expanded-distance formula reduces intermediates but can suffer cancellation; chunked direct subtraction is a useful reference. A compiled loop can be preferable for sequential/irregular algorithms. No JIT dependency is installed here; unsupported compilation behavior is implementation-specific, not always a silent fallback.

8 · Vectorized backward and the four-point scalar-to-array bridge

For affine-ReLU-affine CE, use dZ2=(P-one_hot)/N, dW2=H.T@dZ2, db2=sum(dZ2), dH=dZ2@W2.T, dZ1=dH*(Z1>0), dW1=X.T@dZ1, db1=sum(dZ1). An affine layer's input and weight pullbacks have roughly two forward matmuls' cost, but omitting dX, sparse operations, checkpoint recomputation and other operators change the ratio. “Backward at most2× forward” is not a general bound.

The complete vectorized DL012 network below preserves all three tanh layers and sum squared error. Initialize in neuron order with the same Python generator so scalar/array parameters can be copied exactly. A comparison should include all41 gradients and multiple updates, not merely similar final accuracy.

def init_scalar_order(seed=0):
    r=random.Random(seed);params=[]
    for nin,nout in [(3,4),(4,4),(4,1)]:
        W=np.array([[r.uniform(-1,1) for _ in range(nin)] for _ in range(nout)]).T
        params.append((W,np.zeros(nout)))
    return params
 
def vector_eval(X,y,params):
    inputs=[];acts=[];h=X
    for W,b in params:
        inputs.append(h);h=np.tanh(h@W+b);acts.append(h)
    loss=float(((h-y)**2).sum());dh=2*(h-y);grads=[]
    for W_b,inp,act in zip(reversed(params),reversed(inputs),reversed(acts)):
        W,b=W_b;dz=dh*(1-act*act)
        grads.append((inp.T@dz,dz.sum(0)));dh=dz@W.T
    return loss,list(reversed(grads)),h
 
VX=np.array([[2.,3.,-1.],[3.,-1.,.5],[.5,1.,1.],[1.,1.,-1.]])
VY=np.array([[1.],[-1.],[-1.],[1.]])
vp=init_scalar_order();initial,vg,_=vector_eval(VX,VY,vp)
errors=[]
for li,pair in enumerate(vp):
    for pi,value in enumerate(pair):
        for ix in np.ndindex(value.shape):
            old=value[ix];value[ix]=old+1e-5;plus=vector_eval(VX,VY,vp)[0]
            value[ix]=old-1e-5;minus=vector_eval(VX,VY,vp)[0];value[ix]=old
            errors.append(abs((plus-minus)/2e-5-vg[li][pi][ix]))
assert max(errors)<1e-7
for _ in range(200):
    vl,vg,_=vector_eval(VX,VY,vp)
    vp=[(W-.05*dW,b-.05*db) for (W,b),(dW,db) in zip(vp,vg)]
vl,_,pred=vector_eval(VX,VY,vp)
assert vl<initial and np.array_equal(pred>0,VY>0)
print('vectorized four-point training',initial,vl,'gradient error',max(errors))

The separate review harness compares this exact program with DL012's exact scalar engine at copied parameters and records CPU timings. Those four points have three features; calling them “four-point moons” was incorrect. No MNIST dataset or GPU run is hidden in this comparison.

9 · Multi-head attention and the einsum alternative

For Q,K,V shaped (B,T,H,D), compute scores (B,H,T,T), normalize over key positions, then return (B,T,H,Dv). Do not reuse a2D softmax fixed to axis1: that would normalize over heads for this tensor.

Q=rng.normal(size=(2,3,2,4));K=rng.normal(size=(2,5,2,4));V=rng.normal(size=(2,5,2,6))
s=np.einsum('bshd,bthd->bhst',Q,K)/2
u=s-s.max(-1,keepdims=True);attn=np.exp(u);attn/=attn.sum(-1,keepdims=True)
o=np.einsum('bhst,bthv->bshv',attn,V)
ref_s=Q.transpose(0,2,1,3)@K.transpose(0,2,3,1)/2
ref_o=(attn@V.transpose(0,2,1,3)).transpose(0,2,1,3)
assert np.allclose(s,ref_s) and np.allclose(o,ref_o)
assert np.allclose(attn.sum(-1),1)
wrong=np.exp(u);wrong/=wrong.sum(1,keepdims=True)
assert not np.allclose(wrong.sum(-1),1)

These are unmasked cross-attention dimensions to make unequal query/key lengths explicit. Causal self-attention adds a mask before softmax and must handle fully masked rows explicitly. This fixture teaches tensor contraction, not a production attention kernel.

10 · Failure cases and design choices

Profile before replacing: a per-example feature helper can dominate, but a loop over a handful of layers may be irrelevant. Rank semantics: dot and matmul differ above2D; test expected shape as well as values. Einsum performance: default dispatch/path is not guaranteed to match BLAS; cache a justified path only when planning cost matters. Layout: compare the same mathematical reduction and record strides; a view's existence does not prove a slowdown. Temporary allocation: count expanded shapes before writing all-pairs broadcasts.

Choose broadcasting for small readable elementwise expressions; matmul/einsum for genuine contractions; chunking for memory limits; compiled loops for irregular/sequential kernels when an existing supported compiler is available. Each requires output parity and measured time/memory. Neither vectorization nor compilation is proof of correctness or universal speed.

11 · Retention scaffold — recall and stretch answers

1. Why bulk operations can win: fewer interpreter dispatches, efficient numeric storage, SIMD and cache reuse where applicable. The multiplier is measured, not1000× by definition.
2. Matmul notation: mk,kn->mn; k disappears and is summed. Dot is i,i->.
3. FLOPs: conventional12,845,056 for (64,784)@(784,128); exact scalar operations subtract64*128 additions from that estimate.
4. Batching: more independent work and amortized overhead can improve utilization, while memory and latency may worsen. No GPU result was measured here.
5. Sequential example: sampled token t+1 depends on sampled token t; teacher-forced training differs. Nonlinear recurrent state is generally sequential, with independent batch axes still available.
6. Stretch: §8 gives the complete3→4→4→1 NumPy sum-loss network, all-coordinate finite differences and200 updates. Compare against DL012 at identical weights and report observed time; no promised100× factor.

In your own words, explain why a broadcast expression can use more memory than a loop. Revisit broadcasting and the softmax axis before the next framework lesson. The sticky rule is measure and validate the hot loop, not “never write a Python loop.”

Sources and execution boundary

NumPy2.2 einsum, matmul, dot, copies/views, and Attention Is All You Need. Primary passages support API/contraction contracts; FLOP/intensity arithmetic and negative controls are separate local evidence. CPU timings are process-specific, and no GPU, compiler or MNIST training benchmark is claimed.