Part III — Debugging Interactive and Numerical AI

Same error, two diseases

The data is honest now (Chapter 13): clean split, quarantined columns, loader contracts. Training starts — and dies on the first batch:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The engineer moves the model to CUDA. Next run, a sibling error:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x128 and 256x64)

Both are RuntimeErrors from the same training step. The team debates “a device/config problem” vs. “a layer-dimension problem” — and both sound right, because the traceback points at the matmul line in both cases while the actual handoff broke frames earlier.

OBSERVATION: identical training command → device-mismatch RuntimeError on run 1; shape-mismatch RuntimeError on run 2 after the device patch. HYPOTHESIS H1 (wrong-dimension): a tensor carries an unexpected rank or axis size across a handoff (batch-first vs. sequence-first, flattened vs. unflattened, vocab vs. hidden). HYPOTHESIS H2 (wrong-device/dtype): shapes agree but placement or precision disagree (CPU vs. CUDA, float32 vs. long indices, float16 overflow) at the same handoff. INFERENCE: none yet — the error text names the crash site, not the handoff that produced the bad tensor. Only shape/dtype/device assertions at each handoff separate H1 from H2.

This chapter’s question: when tensor code crashes at a line innocent of the defect, what ordered probe names the guilty handoff?

Why reading the error line fails

The obvious move — “fix the line the traceback names” — fails because tensor errors surface downstream of their cause. Three habits misread them:

1. Treating the crash line as the defect line. matmul in the attention block raises, but the wrong axis was introduced three functions earlier by a reshape/permute/squeeze that silently succeeded. Broadcasting makes it worse: (32, 1, 128) multiplies happily where (32, 128) was intended — until a later op disagrees.
2. Patching devices/dtypes at the crash. .to(device) sprinkled at the error line fixes this run and hides the handoff missing its contract — the next tensor through the same gap crashes identically. Device and dtype are handoff properties, not crash-site bandages.
3. Guessing axes from memory. “Batch is dim 0, surely.” Surely fails against (seq, batch, hidden) convention layers, channels-first images, and padding masks shaped (batch, 1, seq) vs. (batch, seq). Memory is not a shape system; assertions are.

OPINION: every tensor handoff without an asserted (shape, dtype, device) triple is an untyped function boundary. Chapter 8’s contracts were invented for exactly this gap — here they are executable.

The first systematic study of deep-learning defects supports the “surface downstream” claim directly. Zhang and colleagues analyzed TensorFlow bugs mined from Stack Overflow and GitHub and found that tensor-shape mismatches and type/precision errors are among the most common root causes, and that a recurring difficulty is that the symptom appears far from the defect and is hard to reproduce (Zhang et al., 2018). Humbatova and colleagues’ taxonomy (Chapter 4) makes the same point with its large “Tensors & Inputs” category.

The mental model: a tensor is three facts — shape, dtype, device — and every function boundary must re-establish all three. Rank errors, axis swaps, dtype mismatches (long expected, float received in embedding indices), and device splits are four names for one event: a handoff whose contract was assumed, not asserted.

The method: the shape-assertion probe

Ordered probe — read the error triple, then walk upstream asserting at handoffs, one fix per run:

1. Read the full triple, not just the message. For the crashing op, record expected vs. actual (shape, dtype, device) of each operand. RuntimeError text usually gives two of the three — measure the third yourself (t.shape, t.dtype, t.device).
2. Walk upstream to the first handoff where the triple diverges from intent. Insert temporary assertions at each boundary (loader → collate → embedding → encoder → head): assert x.shape == (B, S, H), assert x.dtype == torch.long, assert x.device == device. The first failing assertion upstream is the conviction point; the crash line is its symptom.
3. Run the discriminating probe: one handoff pinned per run. Prediction if H1 (wrong-dim): pinning the axis order at the suspect handoff (permute/reshape corrected, assertion green) moves the crash downstream or clears it, while device/dtype pins change nothing. Prediction if H2 (wrong-device/dtype): a .to()/.to(dtype) at the handoff (not the crash) clears it, while reshape edits change nothing. Never apply both pins between runs.
4. Promote the probe to a contract. The temporary assertion that caught the defect stays — at the handoff, in the committed code, running every batch (Chapter 8’s tripwire, tensor edition).

    flowchart TD
    C["crash op: record both operands' (shape, dtype, device) + the verbatim error"] --> W["insert handoff() prints at each boundary: loader -> collate -> embed -> encoder -> head"]
    W --> F["run once, no fixes: find the first boundary where the triple diverges from intent"]
    F --> P{"which fact diverged?"}
    P -->|"axis / rank"| H1["pin the axis at that handoff (permute / reshape); device and dtype pins change nothing"]
    P -->|"device / dtype"| H2["pin .to() / .to(dtype) at that handoff; reshape edits change nothing"]
    H1 --> R{"crash cleared, or moved downstream?"}
    H2 --> R
    R -->|"moved"| W
    R -->|"cleared"| PR["promote the catching assertion to a permanent handoff contract"]
  

# handoff contract probe (temporary while debugging, permanent once convicted)
def handoff(tag, t, *, shape=None, dtype=None, device=None):
    # MEASUREMENT at the boundary: shape + dtype + device, every batch
    info = f"{tag}: shape={tuple(t.shape)} dtype={t.dtype} device={t.device}"
    print(info)
    if shape is not None:
        assert tuple(t.shape) == shape, f"{info} expected shape={shape}"
    if dtype is not None:
        assert t.dtype == dtype, f"{info} expected dtype={dtype}"
    if device is not None:
        assert str(t.device) == str(device), f"{info} expected device={device}"
    return t

# usage at each boundary:
x = handoff("after-collate", batch_ids, shape=(32, 128), dtype="torch.int64", device="cuda:0")

OBSERVATION (constructed illustration, not a measured run): crash at matmul (32x128 vs 256x64); upstream handoff("after-embed") showed (32, 128, 256) where (32, 256, 128) was intended — a missing permute two frames above the crash. Device pin changed nothing; axis pin cleared the run. UPDATED BELIEF: H1 supported — wrong-dimension at the embedding→encoder handoff; H2 exonerated for this crash (device was uniformly cuda:0 at every assertion). INFERENCE: the fix is the permute plus its permanent assertion — not a reshape at the crash line, which would have silently transposed semantics.

Research lineage: the probe can be static, and the axes can have names

Shape errors can be caught before the run. Jhoo and colleagues’ PyTea statically traces every execution path in PyTorch training code, collects the shape constraints each tensor operation imposes, and checks whether they are jointly satisfiable; it found real shape errors in the official PyTorch repository and in Stack Overflow snippets, each in seconds (Jhoo et al., 2022). The upstream survey in this chapter is the manual, dynamic form of the same analysis — run PyTea first where you can, and the survey becomes confirmation rather than discovery.

The root cause is positional axes. Rush and Chiang argue that indexing tensor dimensions by position — dim=0, dim=1 — is the notational choice that makes axis-swap bugs so easy to write and so hard to read, and propose named tensor notation where every axis carries a name and operations bind by name (Rush & Chiang, 2021). Three practical forms of this reach real code. PyTorch’s (experimental) named tensors attach names to a live tensor’s axes. einops (rearrange(x, "b s h -> b h s")) makes every reshape and permutation name its axes at the call site, so an axis swap is a spelling error rather than a silent transpose. And jaxtyping / torchtyping turn the promoted assertion from this chapter’s probe into a type annotation — Float[Tensor, "batch seq hidden"] — checked at the function boundary. All three are Chapter 8’s contract with the axis names written down; they turn “guessing axes from memory” into an error you can see before the run.

Lab 14: shape-assertion probe separating H1/H2

PROPOSED, not executed: no author-measured results are reported. The evidence this chapter requires is the reader’s own handoff table.

Setup. Take any tensor crash (or inject both: x.permute(0,2,1) removed for an H1 run; one operand left on CPU for an H2 run). Freeze the batch, seed, and input sample — they are controlled variables; the pinned handoff is the independent variable.

Task.

1. Write H1 (wrong-dim at handoff N) and H2 (wrong-device/dtype at handoff N) with distinct predictions before editing: which pin clears the crash, and what the upstream assertion will print.
2. Record the crash triple (both operands’ shape/dtype/device, verbatim error) then insert handoff() prints at each upstream boundary and run once — no fixes yet. The first divergent handoff is the suspect.
3. Apply exactly one pin (axis fix or device/dtype fix at the handoff) and re-run. Record OBSERVATION (green / moved crash with new triple / unchanged) and UPDATED BELIEF. If the crash moves downstream, repeat the walk — one handoff per run.
4. Tensor ops are deterministic given seeds here, but cuDNN/GPU ordering can vary — confirm the final green twice.

Success criterion. The named handoff, the triple that diverged, the single pin that cleared it, and the permanent assertion committed at that boundary. A crash-line patch without the upstream triple is explicitly not completion.

Companion tool: Tensor Shape/Type Inspector

What it accepts: the crash triple, the per-handoff (shape, dtype, device) survey table, the suspect handoff name, and the probe outcomes. What it performs: it enforces walk order (crash triple → upstream survey → one pin per run), refuses a verdict while any handoff row is UNKNOWN, checks probe outcomes against pre-written predictions, and records the promoted assertion with its location. What it can establish: which handoff introduced the divergence and which class (dimension vs. device/dtype) it belongs to — under the examined batch and revision only. What it cannot establish: data honesty (Chapter 13), training dynamics (Chapter 15), or semantic correctness — a green triple can still carry transposed meaning (right shape, wrong axis semantics). It never treats crash-line agreement, a single green batch, or dtype-cast silence as diagnosis. How its output changes your next action: an H1 conviction routes to the axis fix + permanent assertion at the handoff; an H2 conviction routes to device/dtype policy at the handoff (factory functions, already-placed batches); a green survey with a still-crashing run routes to Chapters 13/15 with the tensor layer exonerated in writing.

Paper form, sufficient for this chapter:

Crash op: ___   Err verbatim: ___
Operand A (shape/dtype/device): ___ / ___ / ___   Operand B: ___ / ___ / ___
Handoff survey (boundary → triple → matches intent? Y/N): ___
Suspect handoff: ___   FORECAST H1: ___   FORECAST H2: ___
Pin applied (one): ___   Result: cleared / moved (new triple ___) / unchanged
PROMOTED ASSERTION (file:line): ___

Where a software implementation does not yet exist in the reader’s stack, this record is the tool. The handoff-assertion discipline precedes any automation.

Reusable procedure: every tensor crash gets this

1. Transcribe the triple — both operands, all three facts, verbatim error.
2. Survey upstream — prints at every handoff, no fixes, first divergence wins.
3. Pin one handoff per run — axis or device/dtype, predictions pre-written.
4. Chase moved crashes — each new crash is a new survey from its site upward.
5. Promote the assertion — the catching probe becomes the permanent contract.

Failure modes

• Crash-line surgery. Reshaping at the matmul to make the error stop. The numbers fit; the semantics transposed silently — a wrong model that runs.
• .to(device) confetti. Device casts scattered at crash sites instead of one placement policy at the batch factory. Each new tensor re-opens the gap.
• Silent broadcast reliance. Shapes that “happen to multiply” via broadcasting with an unexamined (…, 1, …) axis. Broadcasting is implicit reshaping without an assertion — audit every broadcast pair.
• Dtype telepathy. Assuming embeddings receive long and losses receive float. Mixed precision and index tensors punish assumptions; assert both.
• Squeeze/unsqueeze drift. squeeze() removing a size-1 batch dim on the last batch, unsqueeze added for one caller breaking another. Rank-polymorphic code needs shape contracts most.
• Single-batch green. One batch passing treated as handoff health. Variable-length batches, last-batch remainders, and empty edge batches each deserve their own triple row.

Limits, per contract: one survey convicts one handoff under one batch/revision; it does not certify data honesty, training dynamics, or axis semantics beyond the asserted triple. UNKNOWN where any handoff row is unmeasured.

References

• Yuhao Zhang, Yifan Chen, Shing-Chi Cheung, Yingfei Xiong, and Lu Zhang. An Empirical Study on TensorFlow Program Bugs. Proceedings of the 27th ACM SIGSOFT International Symposium on Software Testing and Analysis (ISSTA), 2018, pp. 129–140. https://doi.org/10.1145/3213846.3213866
• Ho Young Jhoo, Sehoon Kim, Woosung Song, Kyuyeon Park, DongKwon Lee, and Kwangkeun Yi. A Static Analyzer for Detecting Tensor Shape Errors in Deep Neural Network Training Code. Proceedings of the ACM/IEEE 44th International Conference on Software Engineering: Companion Proceedings (ICSE), 2022. https://doi.org/10.1145/3510454.3528638 · full technical version: https://arxiv.org/abs/2112.09037
• Alexander Rush and David Chiang. Named Tensor Notation. arXiv:2102.13196, 2021. https://arxiv.org/abs/2102.13196

Debugging Checklist

• Crash triple (both operands × shape/dtype/device) transcribed verbatim?
• Upstream handoff survey run with prints only — first divergence named?
• H1/H2 with distinct pin predictions written before editing?
• Exactly one pin per run; moved crashes re-surveyed from the new site?
• Final green confirmed twice (GPU-ordering caution)?
• Catching assertion promoted to permanent handoff contract?
• No crash-line reshape or scattered .to() accepted as the fix?

What This Chapter Established

• Tensors as (shape, dtype, device) triples with contract-bearing handoffs; crash sites as symptoms, first-divergent handoffs as defects.
• The shape-assertion probe (triple → upstream survey → one pin per run → promote), demonstrated on the 32x128 vs 256x64 missing-permute case separated from the cuda:0 vs cpu case — constructed illustration, no measured runs claimed.
• Lab 14 as a proposed handoff record the reader executes; the Tensor Shape/Type Inspector contract (accepts/performs/can-establish/cannot-establish/next-action).
• What was NOT proved: data honesty, training-dynamics health, or semantic axis correctness beyond the asserted shapes.
• Research grounding: shape/type bugs are a top DL defect class that surfaces far from its cause (Zhang et al.); the upstream survey has a static counterpart (PyTea); the root cause is positional axis indexing, and named tensors / jaxtyping turn the promoted assertion into a checked signature (Rush & Chiang).
• Forward link: every handoff is green, batches flow, loss computes — and loss does nothing. The machinery runs; learning does not. That silence is next.

Next

Shapes agree, dtypes agree, devices agree — and the loss curve is flat as glass. Or NaN by epoch two. Or oscillating without descent. The tensors are correct containers carrying a training process that is failing, or a logging process that is lying about one. The next chapter triages training pathology: learning rate, data, loss — and the instrumentation itself as a suspect.