Use with AI · Chapter Application

Audit PyTorch Tensor Shapes and Broadcasting Bugs

Trace tensor dimensions through a real PyTorch codebase and find reshape, view, permute, squeeze, concatenation and silent broadcasting failures before changing model architecture.

Chapter Application PyTorch: Zero to Hero PyTorch tensor and model code Intermediate
Read the source chapter → View the learning path →

How to use this

  1. Open a repository-aware AI assistant.
  2. Give it access to the repository or files you want reviewed.
  3. Copy the prompt below and run it unchanged first.
  4. Use the evidence it finds to decide what to inspect or change next.
PromptCopy and run against your own project
You are reviewing a PyTorch repository for tensor-shape and broadcasting defects.

Do not begin by rewriting the model or adding defensive reshapes. First reconstruct the dimensional contract the code is trying to implement.

Inspect the relevant data pipeline, model forward path, loss computation and any preprocessing/postprocessing that changes tensor dimensions.

Work in this order:

1. Identify the important tensors crossing subsystem boundaries.
   For each one, record:
   - semantic meaning of every dimension
   - expected shape
   - actual shape where it can be inferred
   - dtype and device when relevant

2. Trace every shape-changing operation, including:
   - reshape
   - view
   - flatten
   - squeeze / unsqueeze
   - transpose / permute
   - stack / cat
   - indexing and slicing
   - pooling
   - broadcasting arithmetic

3. Look specifically for high-risk failures:
   - batch and sequence dimensions swapped
   - accidental removal of batch dimension with squeeze
   - view() after a non-contiguous permutation
   - reshape that preserves element count but changes semantic layout
   - silent broadcasting that makes incompatible tensors appear valid
   - concatenation or stacking along the wrong axis
   - hard-coded dimensions that fail for different batch/sequence/image sizes
   - loss targets whose shape is accepted but semantically wrong

4. Separate findings into:
   - confirmed defect
   - likely defect
   - fragile dimensional assumption
   - valid but confusing code

5. For every confirmed or likely defect, provide:
   - exact file/function
   - input shape
   - transformation
   - resulting shape
   - why it is wrong
   - smallest safe correction
   - a runtime assertion or test that would prove the correction

Do not recommend broad refactors unless the evidence shows the current dimensional contract cannot be made reliable locally.

Output:

## Shape map
A compact table of the important tensors and their dimensional meaning.

## Findings
Ranked by severity, with repository evidence.

## Minimal corrections
Only changes justified by the findings.

## Verification
Concrete assertions/tests using representative and edge-case shapes.

If the shape flow is correct, say so explicitly and identify the assumptions that should remain documented or tested.