import torch
import torch.nn.functional as F
torch.manual_seed(0)<torch._C.Generator at 0x1117df5d0>The five posts before this one keep circling one idea: shape is only the first contract. The packed Q/K split that decodes the wrong layout, the loss that broadcasts into an outer product, the distance formula that returns negative distances — in each case the shape assertion passes and the error survives. But shape isn’t the only contract a tensor program has to satisfy. It also has to get the layout right, the coordinate convention right, the numerics right, the device placement right, and the synchronization behaviour right — and violating any of those produces code that runs, returns something plausible, and is wrong or unusable.
This post is where you find out whether that has sunk in. Ten small problems. What makes them worth your time isn’t the solutions — most are two or three lines — it’s the failure mode each one is built around. Some have a tempting wrong answer that returns the correct shape and raises nothing. Others are semantically correct but unusable in a real hot path because of what they cost. Both count.
Read each problem, close the page, write your version. Then open the solution. If you want the reasoning behind any of them, each links back to the post that develops it: strides and views, broadcasting and einsum, gather and indexing, numerical stability, and the GPU cost model.
1. Dot product without torch.dot
a, b = torch.randn(16), torch.randn(16)
# produce the scalar dot product
NoteSolution
a, b = torch.randn(16), torch.randn(16)
c = (a * b).sum(dim=-1)
torch.testing.assert_close(c, torch.dot(a, b))
print("elementwise-multiply-then-reduce matches torch.dot ✓")elementwise-multiply-then-reduce matches torch.dot ✓Why dim=-1 rather than a bare .sum(): it generalizes for free. Batched dot products over (B, D) → (B,) are the same line, provided the leading axes correspond and the last axis is the one being contracted. PyTorch also has a purpose-built version: torch.linalg.vecdot(a, b, dim=-1).
a @ b also works here — but only because both operands happen to be 1-D. Watch what it does when a batch axis appears:
for B, D in [(4, 16), (8, 8)]:
x, y = torch.randn(B, D), torch.randn(B, D)
try:
print(f"B={B}, D={D}: a @ b -> {tuple((x @ y).shape)}"
+ (" <-- valid matmul, WRONG semantics" if B == D else ""))
except RuntimeError:
print(f"B={B}, D={D}: a @ b raises (shapes don't chain) <-- the lucky case")B=4, D=16: a @ b raises (shapes don't chain) <-- the lucky case
B=8, D=8: a @ b -> (8, 8) <-- valid matmul, WRONG semanticsThe trap: when B != D the mismatch raises and you fix it in ten seconds. When B == D — which is easy to hit when both dimensions use common power-of-two sizes — a @ b becomes a perfectly valid matrix multiplication, returns (B, B), and computes something with nothing to do with the row-wise dot products you asked for. The error message you were relying on doesn’t come.
2. Split an interleaved Q/K feature axis
def get_qk(x, nh, hd):
"""x: (B, S, D) where D = nh * 2 * hd.
The feature axis is packed INTERLEAVED: [h0_q, h0_k, h1_q, h1_k, ...]
Return q, k each of shape (B, nh, S, hd)."""
NoteSolution
def get_qk(x, nh, hd):
if x.ndim != 3:
raise ValueError(f"expected (B, S, D), got {tuple(x.shape)}")
if x.shape[-1] != 2 * nh * hd:
raise ValueError(f"expected D = 2*nh*hd = {2*nh*hd}, got {x.shape[-1]}")
packed = x.unflatten(-1, (nh, 2, hd)) # nh slowest, hd fastest
return (packed[:, :, :, 0].transpose(1, 2),
packed[:, :, :, 1].transpose(1, 2))
B, S, nh, hd = 2, 5, 4, 8
x = torch.randn(B, S, nh * 2 * hd)
# reference straight from the index formula: feature = h*(2*hd) + qk*hd + d
q_ref = torch.stack([x[:, :, h*2*hd : h*2*hd + hd ] for h in range(nh)], dim=1)
k_ref = torch.stack([x[:, :, h*2*hd + hd : h*2*hd + 2*hd] for h in range(nh)], dim=1)
q, k = get_qk(x, nh, hd)
assert torch.equal(q, q_ref) and torch.equal(k, k_ref) # BOTH outputs, not just q
print("unflatten(-1, (nh, 2, hd)) decodes the interleaved packing ✓")
# the tempting wrong answer:
qw, _ = x.chunk(2, dim=-1)
qw = qw.view(B, S, nh, hd).transpose(1, 2)
print(f"chunk(2, -1) returns shape {tuple(qw.shape)} — identical to the correct answer")
print(f"...and matches the reference? {torch.equal(qw, q_ref)}")unflatten(-1, (nh, 2, hd)) decodes the interleaved packing ✓
chunk(2, -1) returns shape (2, 4, 5, 8) — identical to the correct answer
...and matches the reference? FalseThe trap: chunk(2, dim=-1) assumes the layout is [all_Q ‖ all_K]. This tensor is interleaved per head. The wrong version returns the same shape, raises nothing, and quietly feeds your attention scrambled vectors. Real codebases genuinely use both packings — nanoGPT’s CausalSelfAttention block-packs a fused QKV, other kernels interleave — which is why the right first move isn’t writing code, it’s establishing how the producer packed the axis. (External implementations change; check the version you’re actually reading against rather than trusting a description like this one.)
Developed in full: strides, views, and contiguity.
3. Patchify (space-to-depth)
img = torch.randn(N, C, H, W); patch = 2
# target shape: (N, C*patch*patch, H//patch, W//patch)
NoteSolution
N, C, H, W, ps = 2, 3, 8, 8, 2
img = torch.randn(N, C, H, W)
if H % ps or W % ps:
raise ValueError("H and W must be divisible by the patch size")
hp, wp = H // ps, W // ps
out = (img.reshape(N, C, hp, ps, wp, ps) # split H -> (hp, ps) and W -> (wp, ps)
.permute(0, 1, 3, 5, 2, 4) # -> (N, C, ps, ps, hp, wp)
.reshape(N, C*ps*ps, hp, wp)) # fold (C, ps, ps) into channels
assert torch.equal(out, torch.nn.PixelUnshuffle(ps)(img))
assert torch.equal(torch.nn.PixelShuffle(ps)(out), img) # and it round-trips
print(f"{tuple(img.shape)} -> {tuple(out.shape)}, matches PixelUnshuffle and inverts cleanly ✓")(2, 3, 8, 8) -> (2, 12, 4, 4), matches PixelUnshuffle and inverts cleanly ✓Worth knowing the built-in exists: this is nn.PixelUnshuffle, and its inverse PixelShuffle is a common learned-upscaling rearrangement in super-resolution architectures (though plenty of systems use resize-convolution or transposed convolution instead). Being able to write it by hand and knowing not to is the useful combination.
The trap is quieter than in the other problems: the channel ordering that falls out is c·p² + ih·p + iw, which is a convention rather than the convention — a different permute produces a different, equally valid ordering. Nothing tells you which one a set of pretrained weights expects. Same class of problem as puzzle 2, wearing different clothes.
4. Sort one tensor by another tensor’s keys
data = torch.randn(4, 12, 3) # (B, S, D)
keys = torch.randint(0, 256, (4, 12)) # (B, S)
# reorder data's S axis by each row's sort order of keys
NoteSolution
B, S, D = 4, 12, 3
data = torch.randn(B, S, D)
keys = torch.randint(0, 256, (B, S))
idx = torch.argsort(keys, dim=-1)
out = torch.take_along_dim(data, idx.unsqueeze(-1), dim=1) # cleanest
alt = torch.gather(data, 1, idx.unsqueeze(-1).expand(-1, -1, D)) # classic
batch = torch.arange(B, device=data.device).unsqueeze(1) # device-aware!
adv = data[batch, idx] # advanced indexing
assert torch.equal(out, alt) and torch.equal(out, adv)
print("take_along_dim, gather, and advanced indexing all agree ✓")take_along_dim, gather, and advanced indexing all agree ✓The rule for gather, once, so you needn’t look it up again: the output takes the shape of the index, and the index replaces the coordinate only along dim — every other coordinate passes through untouched. That’s the whole reason the index must be the same rank as the input, which is what the unsqueeze(-1).expand(...) is for. take_along_dim accepts the singleton feature index axis and broadcasts it across the feature dimension, so that expansion never appears in your code.
The trap here isn’t a wrong answer — it’s a correct one that can’t be used. The Python double-loop is a perfectly valid specification, and it issues O(B·S) tiny indexing dispatches; in a CUDA hot path that’s disqualifying regardless of the exact kernel count. Write the loop to pin the semantics down, then vectorize, then assert they agree. Both halves.
And note the device= on that arange. A helper index tensor built without it lands on the CPU, and the moment data is on a GPU you get a device mismatch — the kind of bug that never appears in your CPU tests.
Worth pre-empting: if you only need the top few rather than a full ordering, don’t sort — topk returns just the selected values and indices instead of the complete permutation argsort produces, which is what makes it the right API when a small subset is all you need. That pattern is mixture-of-experts routing.
5. Build a coordinate grid
H = W = 4 # square on purpose — see below
# want index_map of shape (H, W, 2) where index_map[i, j] == [i, j]
NoteSolution
H = W = 4
ii, jj = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
index_map = torch.stack([ii, jj], dim=-1)
assert index_map[1, 3].tolist() == [1, 3]
print(f"ij: shape {tuple(index_map.shape)}, index_map[1,3] == {index_map[1,3].tolist()}")
# the same call with the other convention, on a SQUARE grid:
xx, yy = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='xy')
wrong = torch.stack([xx, yy], dim=-1)
print(f"xy: shape {tuple(wrong.shape)}, wrong[1,3] == {wrong[1,3].tolist()}")
print(f"same shape? {index_map.shape == wrong.shape} same values? {torch.equal(index_map, wrong)}")ij: shape (4, 4, 2), index_map[1,3] == [1, 3]
xy: shape (4, 4, 2), wrong[1,3] == [3, 1]
same shape? True same values? FalseThe trap is the argument almost nobody passes, and the square grid is what makes it dangerous. indexing='xy' swaps the first two axes. On a 3×5 grid that changes the shape to (5, 3, 2) and you find out immediately. On a 4×4 grid — square, which power-of-two shapes make common — the shape is identical and only the coordinate semantics are transposed, so nothing tells you. The current default is 'ij', and PyTorch warns that it intends to move to 'xy'; either way, published code should pass indexing= explicitly, because coordinate convention is part of a function’s contract rather than something to inherit from a default.
Then the sting: F.grid_sample wants (x, y) order normalized to [-1, 1] — the opposite convention from the (row, col) grid you just built naturally. The two functions you’d most want to use together disagree.
And once you see that meshgrid is only sugar over two broadcast aranges, it stops being something to memorize:
i = torch.arange(H).view(H, 1).expand(H, W)
j = torch.arange(W).view(1, W).expand(H, W)
assert torch.equal(torch.stack([i, j], -1), index_map)
print("pure broadcasting reproduces meshgrid('ij') ✓")pure broadcasting reproduces meshgrid('ij') ✓6. Masked sum of squares — which is cheaper on a GPU, and why?
data = torch.randn(128); mask = torch.rand(128) < 0.7
a = (data * data * mask.type_as(data)).sum()
b = (data[mask] ** 2).sum()
NoteSolution
At moderate mask density with a cheap downstream reduction, a is the better CUDA starting point — and the reason isn’t arithmetic. data[mask] produces a tensor whose size depends on the data. In ordinary eager CUDA execution, returning that compact tensor needs an exact host-visible output size first, which can introduce synchronization — torch.nonzero, exposing the same data-dependent count, documents it explicitly. That makes the host wait for the relevant preceding device work. (Compiled paths may instead represent supported data-dependent sizes symbolically.) The data-dependent output shape also reduces fusion and graph-capture opportunities — modern torch.compile can represent many such operations using unbacked symbolic sizes, but unsupported control flow or backend limits can still graph-break or fail to compile. Option a is fixed-shape and avoids that compaction entirely. The arithmetic it saves is small; host-visible shape discovery can be expensive. But I haven’t measured the crossover, and it genuinely depends on size, density, how much downstream work compaction avoids, compiler mode, and backend — at extreme sparsity selection can win. Treat this as where to start, then benchmark the actual workload.
But neither is the answer I’d write. Watch what happens when a masked-out slot holds garbage:
data = torch.tensor([1., 2., float('inf'), 4.])
mask = torch.tensor([True, True, False, True])
print("multiply-by-mask:", (data * data * mask.float()).sum().item())
safe = torch.where(mask, data, torch.zeros_like(data)) # sanitize the INPUT first
print("sanitize-then-square:", (safe * safe).sum().item())multiply-by-mask: nan
sanitize-then-square: 21.00 × inf = nan, so multiplying by the mask lets a value you explicitly excluded poison the whole reduction. torch.where selects instead — but it’s eager, evaluating both branches, so where(mask, data*data, 0) gives a finite forward pass and a nan gradient. Sanitizing the input before the dangerous operation is the safe dense formulation shown here — not the only safe form, but the one that keeps option a’s static shape while fixing the gradient.
Two posts meet here: the GPU cost model for the stall, numerical stability for the gradient.
7. Grouped-query attention: repeat KV heads to match Q heads
k = torch.randn(B, n_kv_heads, S, hd) # 2 kv heads
# want (B, n_q_heads, S, hd) # 8 q heads
# query head h must attend to kv head h // (n_q // n_kv)
NoteSolution
B, S, hd = 2, 4, 8
n_q, n_kv = 8, 2
if n_q % n_kv:
raise ValueError("n_q must be divisible by n_kv")
rep = n_q // n_kv # 4
k = torch.randn(B, n_kv, S, hd)
reference = torch.stack([k[:, h // rep] for h in range(n_q)], dim=1)
correct = k.repeat_interleave(rep, dim=1)
wrong = k.repeat(1, rep, 1, 1)
llama_5d = k[:, :, None].expand(B, n_kv, rep, S, hd) # stride-0 view, no copy yet
llama = llama_5d.reshape(B, n_q, S, hd) # what HF's repeat_kv does
assert torch.equal(correct, reference)
assert torch.equal(llama, reference)
print("repeat_interleave and expand+reshape both match the reference ✓")
print(f"repeat(1, rep, 1, 1) returns shape {tuple(wrong.shape)} — identical — "
f"and matches? {torch.equal(wrong, reference)}")
def shares_allocation(a, b): # NOT a.data_ptr() — see note below
return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()
print(f"expand shares k's allocation? {shares_allocation(llama_5d, k)}")
print(f"after reshape? {shares_allocation(llama, k)} <- the copy lands here")repeat_interleave and expand+reshape both match the reference ✓
repeat(1, rep, 1, 1) returns shape (2, 8, 4, 8) — identical — and matches? False
expand shares k's allocation? True
after reshape? False <- the copy lands hereThe first trap is a lovely one: repeat tiles → [kv0, kv1, kv0, kv1, …]; repeat_interleave blocks → [kv0, kv0, kv0, kv0, kv1, …]. Grouped-query attention means the second. Both return (B, 8, S, hd), both raise nothing, and one of them silently scrambles which query heads see which keys.
The second trap caught me while writing this. It’s tempting to call expand + reshape zero-copy, since expand is a stride-0 view. Read the storage lines: the expand aliases, but the reshape merging (n_kv, rep) into one axis materializes here — no single stride can express “advance zero for rep−1 steps, then jump,” so no valid view exists. (That’s for the nontrivial case with n_kv > 1 and rep > 1; degenerate shapes like rep == 1 are view-compatible and don’t copy, which is exactly why this is something to check rather than assume.) Note also the helper: comparing untyped_storage().data_ptr() asks “same allocation?”, where a bare data_ptr() asks “same first element?” — two views into one buffer at different offsets disagree on the second question while sharing storage. So both routes materialize the full four-dimensional tensor. I’d previously written that expand + reshape is therefore “the faster copy, which is why HuggingFace prefers it” — but this experiment doesn’t show that; relative performance is backend- and shape-dependent and would need benchmarking. (As of writing, Transformers’ repeat_kv does use the expand-then-reshape form — but that’s a fact about a moving target, so check the version you depend on.)
The way to genuinely avoid the copy is not to materialize at all: a backend with native GQA support may consume the grouped head mapping without explicitly building repeated K/V tensors. F.scaled_dot_product_attention(..., enable_gqa=True) requests that, but it doesn’t guarantee it — the flag is experimental, restricted to particular CUDA backends and head-count combinations, and PyTorch’s own equivalent formulation is written with repeat_interleave. Verify which backend you actually got rather than assuming the replication vanished.
That “verify storage sharing explicitly rather than inferring it from syntax” habit is from strides — and it’s what corrected this paragraph.
8. Top-p (nucleus) sampling mask
# keep the SMALLEST set of tokens whose cumulative probability >= p
def top_p_mask(logits, p=0.9): ...
NoteSolution
def top_p_mask(logits, p=0.9):
if not 0 < p <= 1:
raise ValueError(f"p must be in (0, 1]; got {p}")
probs = torch.softmax(logits, -1) # assumes finite logits, nonempty vocab axis
sp, si = probs.sort(-1, descending=True)
cum = sp.cumsum(-1)
keep_sorted = (cum - sp) < p # cumulative mass BEFORE this token < p
return torch.zeros_like(keep_sorted).scatter_(-1, si, keep_sorted)
def top_p_naive(logits, p=0.9):
probs = torch.softmax(logits, -1)
sp, si = probs.sort(-1, descending=True)
keep_sorted = sp.cumsum(-1) <= p # no shift — the bug
return torch.zeros_like(keep_sorted).scatter_(-1, si, keep_sorted)
# a confident model: the top token alone holds 95% of the mass, p = 0.9
lg = torch.log(torch.tensor([[0.95, 0.03, 0.02]]))
assert top_p_mask(lg, 0.9).sum().item() == 1
print(f"correct keeps {top_p_mask(lg, 0.9).sum().item()} token(s)")
print(f"naive keeps {top_p_naive(lg, 0.9).sum().item()} token(s) <- empty set")correct keeps 1 token(s)
naive keeps 0 token(s) <- empty setThe trap: cumsum <= p drops the token that crosses the threshold. Usually that’s a harmless off-by-one. But when the distribution is sharply peaked and the top token alone exceeds p, it keeps nothing, and sampling from an empty distribution crashes. Production implementations commonly handle this by shifting the removal mask, as the reference below does, or by enforcing a minimum number of retained tokens.
The (cum - sp) < p formulation — “was the cumulative mass before this token still under p?” — expresses the shift in a single comparison. Rather than ask you to take my word that it matches the standard implementation, here is the comparison, runnable:
def top_p_reference(logits, p=0.9):
"""The widely-used formulation: build a removal mask, then shift it right by one
so the token that crosses the threshold is kept, and always keep the top token."""
probs = torch.softmax(logits, -1)
sp, si = probs.sort(-1, descending=True)
remove = sp.cumsum(-1) > p
remove[..., 1:] = remove[..., :-1].clone()
remove[..., 0] = False
return torch.zeros_like(remove).scatter_(-1, si, ~remove)
torch.manual_seed(0)
mismatches = 0
for _ in range(2000):
lg_ = torch.randn(1, 12) * 3
p_ = float(torch.rand(1) * 0.9 + 0.05)
if not torch.equal(top_p_mask(lg_, p_), top_p_reference(lg_, p_)):
mismatches += 1
assert mismatches == 0
print(f"2000 random (logits, p) pairs: {mismatches} mismatches ✓")2000 random (logits, p) pairs: 0 mismatches ✓9. Ragged → padded, with a mask, no Python loop
lens = torch.tensor([3, 1, 4]) # sequence lengths
flat = torch.arange(8).float() # the sequences, concatenated
# want padded (3, 4) and a boolean mask (3, 4)
NoteSolution
lens = torch.tensor([3, 1, 4])
flat = torch.arange(8, dtype=torch.float32)
B, T = lens.shape[0], 4 # T from a KNOWN capacity, not lens.max() — see below
# --- contract. The structural checks are cheap; the value checks read device data,
# --- so in real code they belong in validation/tests, not a CUDA hot path.
if lens.ndim != 1:
raise ValueError(f"lens must have shape (B,); got {tuple(lens.shape)}")
if lens.dtype not in (torch.int32, torch.int64):
raise TypeError("lens must be integer — float lengths silently round in the mask")
if flat.ndim != 1:
raise ValueError("flat must be the 1-D concatenation of the sequences")
if lens.device != flat.device:
raise ValueError("lens and flat must be on the same device")
if bool((lens < 0).any()) or bool((lens > T).any()): # syncs on CUDA
raise ValueError(f"every length must satisfy 0 <= len <= T={T}")
if int(lens.sum()) != flat.numel(): # syncs on CUDA
raise ValueError("sum(lens) must equal flat.numel()")
positions = torch.arange(T, device=lens.device)
mask = positions[None, :] < lens[:, None] # (B, T) — the whole trick
padded = flat.new_zeros((B, T)) # inherits device and dtype
padded.masked_scatter_(mask, flat)
print("padded:\n", padded.int().tolist())
print("mask:\n", mask.int().tolist())padded:
[[0, 1, 2, 0], [3, 0, 0, 0], [4, 5, 6, 7]]
mask:
[[1, 1, 1, 0], [1, 0, 0, 0], [1, 1, 1, 1]]The line worth memorizing is arange(T)[None, :] < lens[:, None]. Broadcasting a row of positions against a column of lengths is the padding mask. You’ll write this, or a close variant, in every serving system you ever touch.
The trap most people fall into is looping over the batch: correct, and O(B) tiny dispatches sitting in a data path.
The traps the contract catches are worth seeing, because they fail silently. A float lens of 2.5 compares as three positions and quietly invents a token. And a length larger than T makes the mask simply stop at the capacity: mask.sum() no longer equals flat.numel(), masked_scatter_ fills what it can, and you get a quietly truncated batch with no exception at all. (A negative length, by contrast, does raise — the loud failure is the lucky one.) Hence the sum(lens) == flat.numel() check, which turns the silent case into a real error.
The trap I fell into is subtler, and it’s worth flagging because it violates this series’ own checklist. My first version wrote T = int(lens.max()) and torch.zeros(B, T). Both are quiet mistakes on a GPU: int(...) on a device tensor forces a host synchronization — precisely what the GPU cost model post spends its length warning about — and torch.zeros without device= builds a CPU tensor. The version above takes T from a known capacity and uses new_zeros to inherit device and dtype. If you genuinely must derive T from the data in eager CUDA code, that sync is unavoidable, which is why production systems bucket sequence lengths, pad to fixed capacities, or use nested/jagged tensors instead.
The follow-on: this is exactly the mask that feeds attention — and if a zero-length row is represented by setting every logit to -inf, a naive softmax returns nan across that whole row unless the implementation defines an explicit empty-row policy. Numerical stability works through the policy choices for that.
10. Batched pairwise distances without the memory blow-up
A = torch.randn(B, N, D); Bt = torch.randn(B, M, D)
# want (B, N, M) squared distances — the naive form OOMs
NoteSolution
Bb, N, M, D = 4, 200, 300, 64
A = torch.randn(Bb, N, D)
Bt = torch.randn(Bb, M, D)
naive = ((A[:, :, None, :] - Bt[:, None, :, :]) ** 2).sum(-1) # (B,N,M,D) intermediate!
raw = ((A**2).sum(-1)[:, :, None]
+ (Bt**2).sum(-1)[:, None, :]
- 2 * torch.bmm(A, Bt.transpose(1, 2))) # (B,N,M) only
# Diagnostic host reads — fine for inspection and tests, but these .item() calls
# would synchronize inside a CUDA hot loop. Keep them out of the fast path.
print(f"minimum raw squared distance: {raw.min().item():.4f} "
f"(negatives: {(raw < 0).sum().item()})")
smart = raw.clamp_min(0) # ...clamp only AFTER looking at the error scale
torch.testing.assert_close(naive, smart, atol=1e-3, rtol=1e-3)
print(f"naive pairwise-difference tensor: {Bb*N*M*D*4/2**20:6.1f} MiB")
print(f"one B x N x M pairwise tensor : {Bb*N*M*4/2**20:6.1f} MiB ({D}x fewer elements)")
print("(agreeing within the stated floating-point tolerance, not bit-for-bit)")minimum raw squared distance: 47.3572 (negatives: 0)
naive pairwise-difference tensor: 58.6 MiB
one B x N x M pairwise tensor : 0.9 MiB (64x fewer elements)
(agreeing within the stated floating-point tolerance, not bit-for-bit)Expanding ‖a−b‖² = ‖a‖² + ‖b‖² − 2a·b replaces a (B,N,M,D) broadcast with a (B,N,M) matmul. The dominant pairwise intermediate loses the feature dimension and shrinks by a factor of D; note that peak memory isn’t exactly one such tensor, since the expression also builds the two norm reductions, the matmul output, and the broadcast sums. At realistic sizes it’s still the difference between an out-of-memory error and a working kernel.
The trap is the clamp_min(0), and it isn’t cosmetic. That expansion subtracts two large, nearly-equal quantities, so in fp32 it can return negative squared distances through catastrophic cancellation — and sqrt of those is nan. Note the printed minimum above: this particular random, roughly-centered data doesn’t trigger it, which is precisely why the habit is to look rather than to clamp reflexively. Move the points far from the origin, as the numerical stability post does, and the negatives appear immediately. Worse, without the sqrt there’s no nan at all: just corrupted magnitudes that never announce themselves. And know what the clamp does and doesn’t buy: it restores the invariant (non-negative, sqrt-safe), not the accuracy — a genuinely small separation between two distant points can come back as a large negative number and clamp to zero, i.e. “these points are identical.”
If near-duplicate geometry is what you actually care about, compute directly (chunked if memory demands) or in higher precision. torch.cdist is the maintained option — note it returns Euclidean distances rather than squared ones — but check which compute_mode you’re getting, since its default switches to this same expanded route once either point count exceeds 25. That whole investigation is numerical stability, and this problem is the one on the list that has genuinely broken production systems.
The checklist
Before calling any of these done:
Each of these problems exposes a contract that shape alone cannot establish. Sometimes the shape is identical and the values are wrong; sometimes the values are right but the layout, device, numerical, or execution contract makes the code unusable. Shape is the first contract your code has to satisfy. It is nowhere near the last, and that’s the whole reason the preceding five posts exist.