2. Tensors: the data structure of deep learning

Why start here

Every number a language model touches lives in a tensor: the weights you load from disk, the token vectors flowing through layers, the attention scores, the KV cache. Every bug you’ll hit while building an engine is, at bottom, a tensor bug: the wrong axis, the wrong layout, the wrong dtype, or the wrong device. An hour spent getting fluent here saves days later.

A tensor is an array with four properties

Run this in your environment (python inside the activated venv):

import torch
x = torch.arange(24.0).reshape(2, 3, 4)
print(x.shape, x.dtype, x.device, x.stride())

torch.Size([2, 3, 4]) torch.float32 cpu (12, 4, 1)

Four properties describe this tensor completely:

• Shape [2, 3, 4]: two blocks, each with three rows of four numbers. In language models the most common shape is [B, T, D]: a Batch of sequences, each with T tokens, each token a vector of D numbers. Our x could be two sentences of three tokens, each token a 4-number vector.
• Dtype float32: how each number is stored. FP32 takes 4 bytes, BF16 and FP16 take 2, INT8 takes 1. Token IDs are integers (int64 by default). The dtype decides memory size, and as Chapter 1 showed, memory size decides decode speed.
• Device cpu: where the data lives and where operations on it run. x.to("cuda") copies it to GPU memory. Operations between tensors on different devices are errors, not silent copies.
• Strides (12, 4, 1): how to find an element in memory. That needs its own section.

Note

Axis meaning lives in your program, not in the tensor. PyTorch doesn’t know that axis 1 is “tokens”. Writing shapes in comments, like # [B, T, D], is the single most effective habit for avoiding bugs in model code. This book does it everywhere.

Memory is flat: strides

RAM is one long line of bytes. A tensor stores its elements in a flat buffer, and the strides say how far to jump in that buffer to move one step along each axis. For our [2, 3, 4] tensor stored row by row, moving one step along the last axis moves 1 element, along the middle axis 4 elements (one row), and along the first 12 (one whole 3×4 block). The element at index (i, j, k) lives at

$$ \text{offset}(i, j, k) = \text{storage_offset} + 12i + 4j + 1k . $$

So x[1, 2, 3] is at 12 + 8 + 3 = 23, the last element. Strides like these, where the last axis moves fastest, are called contiguous or row-major.

Here’s why strides matter: many operations change only the strides, not the data. Transposing swaps two axes by swapping their sizes and strides. No numbers move:

t = x.transpose(1, 2)
print(t.shape, t.stride(), t.is_contiguous())

torch.Size([2, 4, 3]) (12, 1, 4) False

t is a view: another way of looking at the same memory. Views are free, which is why PyTorch uses them so much. The catch is that code which assumes row-major layout will read a view wrongly. A hand-written GPU kernel that computes row * width + col will happily read the wrong numbers from t, and the shapes will all look right. This exact bug shows up when loading checkpoints (Chapter 9) and writing kernels (Part III).

Some operations need contiguous memory. view refuses to reinterpret a non-contiguous tensor; reshape copies when it has to; contiguous() always produces a compact copy if needed:

t.view(2, 12)            # RuntimeError: view size is not compatible with input tensor's size and stride
t.reshape(2, 12)         # works: silently copies
t.contiguous().stride()  # (12, 3, 1): a fresh row-major copy

A copy costs a full read and write of the tensor. In the performance chapters you’ll learn to spot hidden .contiguous() copies in a profile.

The same idea in C++ and Rust: a tensor struct is just a buffer, a shape and strides, and a transpose swaps two entries:

def contiguous_strides(shape):
    strides, step = [], 1
    for size in reversed(shape):
        strides.append(step)
        step *= size
    return tuple(reversed(strides))

def element_offset(index, strides, storage_offset=0):
    return storage_offset + sum(i * s for i, s in zip(index, strides))

// A tensor is a flat buffer plus shape plus strides; element (i, j, ...) is at
// offset + i*stride[0] + j*stride[1] + ...
struct Tensor {
    std::vector<float> data;
    std::vector<size_t> shape, strides;
    size_t offset = 0;

    static std::vector<size_t> contiguous(const std::vector<size_t>& shape) {
        std::vector<size_t> s(shape.size(), 1);
        for (size_t a = shape.size(); a-- > 1;) s[a - 1] = s[a] * shape[a];
        return s;
    }
    Tensor(std::vector<float> d, std::vector<size_t> sh)
        : data(std::move(d)), shape(sh), strides(contiguous(sh)) {}
    float at(const std::vector<size_t>& index) const {
        size_t flat = offset;
        for (size_t a = 0; a < index.size(); ++a) flat += index[a] * strides[a];
        return data[flat];
    }
    Tensor transpose(size_t a, size_t b) const {    // swap sizes and strides; no data moves
        Tensor t = *this;
        std::swap(t.shape[a], t.shape[b]);
        std::swap(t.strides[a], t.strides[b]);
        return t;
    }
};

#![allow(unused)]
fn main() {
/// A view over a flat f32 buffer. Element (i0, i1, ...) lives at
/// offset + i0*strides[0] + i1*strides[1] + ...
#[derive(Clone, Debug)]
pub struct Tensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,
    pub offset: usize,
}

impl Tensor {
    /// Row-major ("contiguous") strides: the last axis moves fastest.
    pub fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
        let mut strides = vec![1; shape.len()];
        for axis in (0..shape.len().saturating_sub(1)).rev() {
            strides[axis] = strides[axis + 1] * shape[axis + 1];
        }
        strides
    }

    pub fn from_vec(data: Vec<f32>, shape: &[usize]) -> Self {
        assert_eq!(data.len(), shape.iter().product::<usize>(), "data length must match the shape");
        Tensor { data, strides: Self::contiguous_strides(shape), shape: shape.to_vec(), offset: 0 }
    }

    pub fn get(&self, index: &[usize]) -> f32 {
        let flat: usize = index.iter().zip(&self.strides).map(|(i, s)| i * s).sum();
        self.data[self.offset + flat]
    }

    /// Swapping two axes swaps their sizes and strides. No data moves.
    pub fn transpose(&self, a: usize, b: usize) -> Self {
        let mut t = self.clone();
        t.shape.swap(a, b);
        t.strides.swap(a, b);
        t
    }

    pub fn is_contiguous(&self) -> bool {
        self.strides == Self::contiguous_strides(&self.shape)
    }
}
}

Elementwise operations and broadcasting

Arithmetic between same-shape tensors works element by element: a + b, a * b, torch.exp(a). More interesting is what happens when shapes differ. Broadcasting lets a smaller tensor be reused across a larger one without copying:

b = torch.tensor([10., 20., 30., 40.])     # shape [4]
print((x + b)[0, 0])                        # the same b added to every 4-vector in x

tensor([10., 21., 32., 43.])

The rule: align shapes from the right. Two sizes are compatible if they’re equal or one of them is 1; a missing axis counts as 1. The result takes the larger size on each axis. So [2, 3, 4] + [4] works (bias added to every token vector), and so does [3, 4] * [3, 1] (each row scaled by its own number):

s = torch.tensor([[1.], [2.], [3.]])       # shape [3, 1]
print(x[0] * s)                             # row i multiplied by s[i]

tensor([[ 0.,  1.,  2.,  3.],
        [ 8., 10., 12., 14.],
        [24., 27., 30., 33.]])

But [2, 3, 4] + [2, 3] fails: aligned from the right, 4 meets 3. If you meant “one number per token”, the scalar tensor must be [2, 3, 1]. Use s.unsqueeze(-1) or s[..., None] to add that axis. You’ll write this rule yourself in the milestone.

Warning

Broadcasting can also succeed when you didn’t intend it. [T, 1] - [T] gives a [T, T] matrix, not a [T] vector. If a loss value or an attention matrix suddenly has an unexpected extra axis, an accidental broadcast is the usual suspect. Assert shapes.

Reductions

A reduction collapses an axis: sum, mean, max, argmax, softmax’s denominator. You name the axis you want to collapse:

print(x.sum(dim=-1))        # sum each 4-vector: [2, 3, 4] -> [2, 3]
print(x.mean(dim=(0, 1)))   # average over batch and tokens: [2, 3, 4] -> [4]

tensor([[ 6., 22., 38.],
        [54., 70., 86.]])
tensor([10., 11., 12., 13.])

keepdim=True keeps the collapsed axis with size 1, which is exactly what you need to broadcast the result back: x - x.mean(-1, keepdim=True) centers every vector. That one line is the first half of LayerNorm (Chapter 6).

Matrix multiplication is many dot products

The dot product of two vectors multiplies them element by element and sums:

$$ [1, 2, 3] \cdot [4, 5, 6] = 1\cdot 4 + 2\cdot 5 + 3\cdot 6 = 32 . $$

A matrix product computes a dot product for every (row of A, column of B) pair:

$$ C_{ij} = \sum_{k} A_{ik} B_{kj}, \qquad [M, K] ;@; [K, N] ;\rightarrow; [M, N]. $$

The shared dimension K is summed away. That’s the shape rule to check first whenever a matmul fails. The work is $M \cdot N \cdot K$ multiply-adds, which we count as $2MNK$ floating-point operations (FLOPs). A transformer spends almost all of its arithmetic here.

def matmul_loops(a, b):
    m, k = a.shape
    k2, n = b.shape
    assert k == k2
    c = torch.zeros(m, n)
    for i in range(m):
        for j in range(n):
            c[i, j] = sum(a[i, p] * b[p, j] for p in range(k))
    return c

// C[M,N] = A[M,K] B[K,N], row-major. The i-k-j order streams rows of B and C sequentially.
inline std::vector<float> matmul(const std::vector<float>& a, const std::vector<float>& b, size_t m, size_t k, size_t n) {
    std::vector<float> c(m * n, 0.f);
    for (size_t i = 0; i < m; ++i)
        for (size_t p = 0; p < k; ++p) {
            float aip = a[i * k + p];
            for (size_t j = 0; j < n; ++j) c[i * n + j] += aip * b[p * n + j];
        }
    return c;
}

#![allow(unused)]
fn main() {
/// C[m][n] = sum_k A[m][k] * B[k][n] for row-major A [M,K] and B [K,N].
/// The i-k-j loop order walks B and C along rows, which keeps memory access sequential.
pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
    let mut c = vec![0.0f32; m * n];
    for i in 0..m {
        for p in 0..k {
            let a_ip = a[i * k + p];
            let b_row = &b[p * n..(p + 1) * n];
            let c_row = &mut c[i * n..(i + 1) * n];
            for (c_ij, b_pj) in c_row.iter_mut().zip(b_row) {
                *c_ij += a_ip * b_pj;
            }
        }
    }
    c
}
}

The Python triple loop is about a million times slower than a @ b. The C++ and Rust versions reorder the loops (i, then k, then j) so the inner loop walks memory sequentially. That’s your first glimpse of a theme that runs through Part III: the same arithmetic, arranged to respect memory, runs far faster.

Batched matmul and the linear layer

With more than two axes, @ treats the leading axes as a batch and multiplies the last two:

print((x @ torch.ones(4, 5)).shape)    # [2, 3, 4] @ [4, 5] -> [2, 3, 5]

torch.Size([2, 3, 5])

One weight matrix transforms every token vector in every sequence. That’s a linear layer, the most common operation in a transformer. One convention matters enormously when loading real weights: PyTorch’s nn.Linear(in, out) stores its weight as [out, in] and computes

$$ y = x,W^{\top} + b . $$

Some checkpoints (GPT-2’s original format, Chapter 9) store [in, out] instead. For square matrices both orientations have the same shape, so a missing transpose passes every shape check and produces garbage. You’ll meet this bug in Chapter 9 and will be glad you know about it.

How big is a tensor?

Bytes = number of elements × bytes per element:

print(torch.tensor(1.0).element_size(), torch.tensor(1.0, dtype=torch.bfloat16).element_size())  # 4 2

Qwen3-0.6B’s embedding table is [151936, 1024]: 155.6 million numbers, which is 311 MB in BF16 and 622 MB in FP32. Get in the habit of doing this multiplication for every big tensor you meet. Memory capacity decides what fits; memory bandwidth decides how fast it runs.

Devices and the GPU

x.to("cuda") copies a tensor into GPU memory and returns the copy; operations on it then run on the GPU. Two facts to keep in mind until Chapter 10 explains them properly:

1. GPU operations are asynchronous. y = x @ x on CUDA returns immediately; the GPU works in the background. Anything that needs the value on the CPU (print(y), y.item(), y.tolist()) waits. Naive timing code therefore measures nothing useful.
2. Copies between CPU and GPU are slow compared with GPU memory itself. Keep data on the device; move only small results back.

Important

On DGX Spark the CPU and GPU share physical memory, but PyTorch still treats them as separate devices with separate allocations. x.to("cuda") is still a copy, and .cpu() still synchronizes.

Build it

pytest tests/test_ch02_tensors.py

Tip

For broadcast_shapes, walk the axes from the right with an index i = 1, 2, ... and treat a missing axis as size 1. For linear, one line is enough: the leading axes broadcast through @ on their own.

Stretch exercises

1. ★ Predict, then check: the shapes of torch.randn(8, 1, 6, 1) + torch.randn(7, 1, 5), torch.randn(2, 3).T @ torch.randn(2, 3), and torch.randn(4, 5)[:, None, :] - torch.randn(4, 5)[None, :, :]. Where: experiments/ch02.py (create it).
2. ★★ Write transpose_strides(shape, strides, a, b) and slice_view(shape, strides, offset, axis, start, step) that return the new (shape, strides, offset) without touching data. Verify against x.transpose and x[:, 1::2]. Where: add both helpers to engine/tensors.py.
3. ★★ Time matmul_loops against @ for 64×64 matrices. How many times slower is it? Use torch.utils.benchmark.Timer. Where: experiments/ch02.py (create it), importing engine.tensors.matmul_loops.
4. ★★★ Implement matmul_loops for any batched shapes [..., M, K] @ [..., K, N] with broadcasting on the leading axes, using your broadcast_shapes. Where: extend matmul_loops in engine/tensors.py.

Check your understanding

1. Why does [2, 3, 4] @ [4, 5] produce [2, 3, 5]?
2. A tensor has shape [3, 2] and strides (1, 3). Is it contiguous? Which operation probably produced it?
3. Why can a missing transpose of a square weight matrix pass every shape check?
4. How many bytes does a [28, 2, 8, 4096, 128] BF16 tensor take? (It’s a KV cache you’ll meet in Chapter 16.)

Going deeper

• BALLM Appendix A, §§A.1-A.5 (pp. 251-271): PyTorch tensors and autograd for newcomers.
• PyTorch documentation, Tensor Views and Broadcasting semantics.
• Edward Yang, PyTorch internals (blog, 2019): how strides, storage and dispatch actually work inside PyTorch.