Module 1 · PyTorch Foundations

Tensors, Shapes, Dtypes, and Broadcasting

Lesson 2 of 32 11 min read Updated

PyTorch models are built from tensor operations. A tensor is a typed, rectangular block of values with a shape, a dtype, a device, and a memory layout. This lesson gives you the working model needed to read tensor code, diagnose shape errors, and choose numeric representations deliberately before you build larger neural networks.

The outcome is practical: you should be able to look at an operation, predict its output shape and dtype, explain when broadcasting applies, and fix common failures without guessing. These skills sit at the beginning of a deep learning course because linear layers, losses, attention mechanisms, image batches, embeddings, and optimizer updates all depend on the same tensor rules.

What a PyTorch Tensor Stores

A tensor object is a view over storage plus metadata. The storage contains the actual values. The metadata tells PyTorch how to interpret those values: shape records the length of each dimension, dtype records the numeric type, device records where operations run, and stride records how many storage steps PyTorch moves when an index changes along each dimension.

Shape is written as a tuple such as (3, 4). The number of dimensions is the rank, exposed as ndim or dim(). A scalar has shape (), a vector might have shape (5,), a matrix might have shape (3, 4), and a mini-batch of images often has shape (batch, channels, height, width).

Dtype controls both meaning and cost. Floating point tensors such as torch.float32 are common for model parameters and activations. Integer tensors such as torch.int64 are common for class labels and token ids. Boolean tensors represent masks. Operations usually require compatible dtypes, and some APIs are strict: classification targets for common loss functions are often integer class indices, while predictions are floating point logits.

Device controls placement. A CPU tensor and a GPU tensor cannot be combined directly in one arithmetic operation. PyTorch will not silently copy one to the other because hidden transfers would make performance and memory behavior unpredictable. You move tensors intentionally with .to(device).

Shape and Layout Internals

Shape alone does not describe memory. A tensor can be contiguous, meaning adjacent logical elements follow the expected row-major storage order, or non-contiguous, meaning it views the same storage with different strides. Transposing a matrix usually creates a view with changed strides rather than copying data. That is efficient, but it affects which reshaping operations are legal.

view() needs a compatible layout because it reinterprets existing storage. reshape() returns a view when it can and makes a copy when it must. permute() and transpose() reorder dimensions by adjusting strides. contiguous() materializes a tensor in standard layout when a later operation needs it.

This distinction matters in neural networks. A convolution might produce an activation tensor that is later flattened before a linear layer. If dimensions are accidentally permuted or a non-contiguous tensor is viewed incorrectly, the model may fail immediately or, worse, feed features in the wrong order.

API Anatomy

Common tensor constructors include torch.tensor for values from Python data, torch.zeros, torch.ones, torch.arange, and torch.empty. Prefer constructors that declare shape and dtype when the tensor is part of a model pipeline. It makes assumptions visible.

Shape inspection uses x.shape, x.ndim, x.numel(), and sometimes x.stride(). Shape changes use unsqueeze to add a length-one dimension, squeeze to remove length-one dimensions, reshape to reorganize elements, and flatten to collapse selected dimensions.

Dtype conversion uses x.to(torch.float32), x.float(), x.long(), or constructor arguments such as dtype=torch.float32. Device and dtype can be changed together with x.to(device=device, dtype=torch.float32).

Broadcasting Rules

Broadcasting lets PyTorch apply elementwise operations to tensors with different shapes when their dimensions are compatible. PyTorch compares shapes from right to left. For each aligned dimension, the sizes must be equal, or one of them must be 1, or one dimension may be missing. Missing leading dimensions behave as if they were length one.

For example, a tensor of shape (2, 3) can be added to a tensor of shape (3,). The vector is treated like one row and reused for each of the two rows. A tensor of shape (4, 1, 8) can combine with (3, 8), producing (4, 3, 8). A tensor of shape (2, 3) cannot combine with (2,) because the rightmost dimensions 3 and 2 conflict.

Broadcasting is a logical expansion. PyTorch does not need to physically copy the smaller tensor for many operations. That makes it powerful and efficient, but it can hide mistakes when a length-one dimension is present accidentally. Always name or comment dimensions in nontrivial code.

Example 1: Inspecting Shape, Dtype, and Strides

This example builds a small matrix, transposes it, and prints the metadata. The values do not change, but the transposed tensor walks through the same storage with different strides.

import torch

x = torch.arange(12, dtype=torch.float32).reshape(3, 4)
y = x.t()

print(x.shape, x.dtype, x.stride(), x.is_contiguous())
print(y.shape, y.dtype, y.stride(), y.is_contiguous())

The expected output shows why shape and layout are separate concepts. The original matrix has shape (3, 4) and standard strides. The transpose has shape (4, 3) but is not contiguous because it is a view over the same storage.

Example 2: Dtype Promotion and Explicit Casting

Mixed numeric operations may promote dtypes. That can be useful, but deep learning code should not rely on accidental promotion across a full training path. Inputs, labels, parameters, and losses should use the dtypes expected by the operation.

import torch

counts = torch.tensor([1, 2, 3], dtype=torch.int64)
scale = torch.tensor(0.5, dtype=torch.float32)
scaled = counts * scale
labels = torch.tensor([0, 2, 1], dtype=torch.int64)

print(scaled)
print(scaled.dtype)
print(labels.dtype)

The product is floating point because multiplying integer counts by a floating scale creates fractional values. The labels remain torch.int64, which is the conventional dtype for class-index targets in many PyTorch classification workflows.

Example 3: Broadcasting Batch Normalization by Hand

Suppose a mini-batch has shape (batch, features). To normalize each feature column, compute per-feature mean and standard deviation along the batch dimension. Keeping dimensions with keepdim=True leaves tensors shaped (1, features), which broadcast cleanly across all rows.

import torch

batch = torch.tensor([[1.0, 10.0, 100.0],
                      [2.0, 20.0, 200.0]])
mean = batch.mean(dim=0, keepdim=True)
centered = batch - mean

print(mean.shape)
print(mean)
print(centered)

The mean has shape (1, 3). Subtracting it from the (2, 3) batch broadcasts the single row over both examples. The centered output is [[-0.5, -5.0, -50.0], [0.5, 5.0, 50.0]].

Example 4: Matrix Multiplication Is Not Elementwise Broadcasting

Broadcasting applies to elementwise operations and also to leading batch dimensions in some linear algebra operations, but matrix multiplication still has strict inner-dimension rules. For a @ b, the last dimension of a must match the second-to-last dimension of b.

import torch

features = torch.ones(2, 3)
weights = torch.arange(12, dtype=torch.float32).reshape(3, 4)
out = features @ weights

print(out.shape)
print(out)

The output has shape (2, 4). Each of the two feature rows multiplies the (3, 4) weight matrix. This is the same shape rule used inside a linear layer before adding bias.

Design Choices and Trade-offs

Use explicit shapes where code crosses a boundary: dataset output, collate function, model forward method, loss calculation, and metric calculation. Assertions are cheap compared with debugging a silent shape mix-up after a long training run.

Choose float32 as the default mental model for training unless a specific mixed-precision workflow is in use. Lower precision can improve throughput and memory use on suitable hardware, but it changes numerical behavior and requires attention to loss scaling and sensitive operations. Higher precision can improve numerical stability for some scientific workloads, but it costs memory and compute.

Use broadcasting for clear per-feature, per-channel, or per-batch operations. Avoid clever broadcasting when named dimensions would be clearer as explicit unsqueeze calls. A short shape comment such as # x: [batch, channels, height, width] often prevents expensive mistakes.

Prefer reshape over view when you do not require view-only behavior. Prefer view when you intentionally want an error if the tensor is not laid out compatibly. That error can reveal an unintended transpose or permute earlier in the pipeline.

Failure Modes and Troubleshooting

Shape mismatch in elementwise operation. Symptom: PyTorch raises an error saying tensor sizes must match at a non-singleton dimension. Cause: broadcasting comparison from the right found conflicting sizes, such as (2, 3) plus (2,). Diagnose by printing both .shape values immediately before the failing operation. Correct by reshaping the smaller tensor to the intended alignment, for example bias.unsqueeze(1) or bias.unsqueeze(0).

Wrong dtype for a loss function. Symptom: a classification loss raises a target dtype error or produces nonsensical behavior after labels were converted to floats. Cause: predictions and targets have different semantic roles. Logits are floating point; class-index targets are integer labels. Diagnose by printing logits.dtype, target.dtype, and representative values. Correct by keeping targets as torch.long when the API expects class indices.

Device mismatch. Symptom: an operation reports tensors on different devices. Cause: model parameters were moved to an accelerator but an input, mask, or newly created tensor stayed on CPU. Diagnose with tensor.device for every operand. Correct by creating new tensors on the existing tensor device, for example torch.zeros_like(x) or torch.arange(n, device=x.device).

Non-contiguous view failure. Symptom: view() fails after transpose() or permute(). Cause: the tensor’s strides cannot represent the requested shape as a simple view. Diagnose with x.is_contiguous() and x.stride(). Correct with x.reshape(...) if a copy is acceptable, or x.contiguous().view(...) if you want to materialize standard layout first.

Performance and Reliability Implications

Tensor metadata choices affect runtime. Unnecessary dtype conversions allocate memory and can block efficient kernels. Repeated CPU-to-GPU transfers inside a training loop can dominate execution time. Accidental broadcasting to a huge result can exhaust memory even when the original tensors were small.

Reliability comes from checking shapes where assumptions change. In a model, validate input rank and feature dimensions at the start of forward when the error message would otherwise appear deep inside a matrix multiply. In data loading, verify that images, labels, masks, and sequence lengths preserve the expected batch dimension. In metrics, detach tensors before storing long-lived results so you do not accidentally retain computation graphs.

Hands-on Lab: Debug a Mini Tensor Pipeline

Prerequisites: Python with PyTorch installed, a terminal, and a clean scratch directory. No dataset or GPU is required.

  1. Create a script named tensor_lab.py with the code below.
  2. Run it once and confirm that all printed shapes match the comments.
  3. Change bias = torch.tensor([0.1, 0.2, 0.3]) to bias = torch.tensor([0.1, 0.2]) and rerun it to observe a broadcasting failure.
  4. Restore the three-element bias, then change labels.long() to labels.float() and observe why label dtype checks are useful.
  5. Cleanup by deleting the scratch script after recording the two failure messages and their fixes.
import torch

x = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
bias = torch.tensor([0.1, 0.2, 0.3])
labels = torch.tensor([0, 1])

assert x.shape == (2, 3)
assert bias.shape == (3,)
assert labels.long().dtype == torch.int64

y = x + bias
z = y.mean(dim=1)

print("x", x.shape, x.dtype)
print("y", y.shape, y.dtype)
print("z", z.shape, z)

Verification: the successful run prints x torch.Size([2, 3]) torch.float32, y torch.Size([2, 3]) torch.float32, and a length-two vector for z. The bias failure proves that broadcasting aligns from the right. The label edit proves that dtype is part of the tensor contract, not cosmetic metadata.

Assessment Exercises

  1. A tensor has shape (8, 1, 32). Another has shape (4, 32). What is the broadcasted output shape for an elementwise add, and why?
  2. You transpose a tensor from shape (5, 7) to (7, 5) and then view(35) fails. Explain the role of strides and give two valid corrections.
  3. A model receives input shape (64, 10) and a weight matrix shape (8, 10). Why does x @ weight fail, and what shape should the weight have for an output with eight features?
  4. Write a three-line check at the start of a forward method that rejects any input that is not a rank-two floating point tensor with 10 features.
  5. Describe one case where broadcasting improves readability and one case where it could hide a bug.

Summary

Tensors are values plus metadata: shape, dtype, device, and layout. Shape determines which operations are legal, dtype determines numeric interpretation and API compatibility, device determines where computation runs, and strides determine how logical indices map to storage. Broadcasting compares dimensions from the right and reuses length-one or missing dimensions for elementwise operations. Mastering these rules makes the later PyTorch topics in this course much easier: every layer, loss, optimizer step, and batch pipeline is built from the same mechanics.

Progress is saved in this browser.