Two tensors, the same shape, and not the same tensor

Attention is a sequence of explicit relationships between represented vectors, and almost every stage can be wrong while producing exactly the shape you predicted. Splitting the embedding across heads is the smallest example: a reshape gives the shape you wanted and not the arrangement of values you meant.

This program composes the chapter’s pieces into one run: the good and bad head split, a round trip that catches the bad one, a hand-written attention compared against PyTorch’s, and a future-token intervention.

Run it

Save this as attention_heads.py and run it with Python 3.11 or later. It needs PyTorch: pip install torch, then python attention_heads.py.

"""Attention: which position is comparing with which, and how to prove it."""
import math

import torch
import torch.nn.functional as F

B, T, E, Nh = 2, 5, 12, 3
Dh = E // Nh


def split_heads(x, num_heads):
    B, T, E = x.shape
    return x.reshape(B, T, num_heads, E // num_heads).transpose(1, 2)


def merge_heads(x):
    B, Nh, T, Dh = x.shape
    return x.transpose(1, 2).reshape(B, T, Nh * Dh)


def manual_attention(q, k, v, allow):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    masked = scores.masked_fill(~allow, float("-inf"))
    weights = torch.softmax(masked, dim=-1)
    return weights @ v, weights


x = torch.arange(B * T * E).reshape(B, T, E)
good = split_heads(x, Nh)
bad = x.reshape(B, Nh, T, Dh)

print(f"good {tuple(good.shape)}  bad {tuple(bad.shape)}  shapes equal: {good.shape == bad.shape}")
print(f"same dtype: {good.dtype == bad.dtype}   same values: {torch.equal(good, bad)}")
print(f"round trip, good: {torch.equal(merge_heads(good), x)}   round trip, bad: {torch.equal(merge_heads(bad), x)}")
print("good head 0 of batch 0:")
print(good[0, 0])
print("bad head 0 of batch 0:")
print(bad[0, 0])

torch.manual_seed(0)
q, k, v = torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh)
allowed = torch.ones(T, T, dtype=torch.bool).tril()
mine, weights = manual_attention(q, k, v, allowed)
theirs = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
print(f"manual vs SDPA, max abs diff: {(mine - theirs).abs().max():.2e}")
print(f"row sums, max error from 1: {(weights.sum(-1) - 1).abs().max():.2e}")
forbidden = weights.masked_select(~allowed)
print(f"forbidden weight max: {forbidden.abs().max():.1e}")

i = 2
v2 = v.clone()
v2[:, :, i + 1:] = torch.randn_like(v2[:, :, i + 1:])
y1 = manual_attention(q, k, v, allowed)[0]
y2 = manual_attention(q, k, v2, allowed)[0]
print(f"max |y1-y2| over positions 0..{i}: {(y1[:, :, :i + 1] - y2[:, :, :i + 1]).abs().max():.2e}")
print(f"max |y1-y2| over positions {i + 1}..: {(y1[:, :, i + 1:] - y2[:, :, i + 1:]).abs().max():.4f}")

"""Attention: which position is comparing with which, and how to prove it."""
import math

import torch
import torch.nn.functional as F

B, T, E, Nh = 2, 5, 12, 3
Dh = E // Nh


def split_heads(x, num_heads):
    B, T, E = x.shape
    return x.reshape(B, T, num_heads, E // num_heads).transpose(1, 2)


def merge_heads(x):
    B, Nh, T, Dh = x.shape
    return x.transpose(1, 2).reshape(B, T, Nh * Dh)


def manual_attention(q, k, v, allow):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    masked = scores.masked_fill(~allow, float("-inf"))
    weights = torch.softmax(masked, dim=-1)
    return weights @ v, weights


x = torch.arange(B * T * E).reshape(B, T, E)
good = split_heads(x, Nh)
bad = x.reshape(B, Nh, T, Dh)

print(f"good {tuple(good.shape)}  bad {tuple(bad.shape)}  shapes equal: {good.shape == bad.shape}")
print(f"same dtype: {good.dtype == bad.dtype}   same values: {torch.equal(good, bad)}")
print(f"round trip, good: {torch.equal(merge_heads(good), x)}   round trip, bad: {torch.equal(merge_heads(bad), x)}")
print("good head 0 of batch 0:")
print(good[0, 0])
print("bad head 0 of batch 0:")
print(bad[0, 0])

torch.manual_seed(0)
q, k, v = torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh)
allowed = torch.ones(T, T, dtype=torch.bool).tril()
mine, weights = manual_attention(q, k, v, allowed)
theirs = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
print(f"manual vs SDPA, max abs diff: {(mine - theirs).abs().max():.2e}")
print(f"row sums, max error from 1: {(weights.sum(-1) - 1).abs().max():.2e}")
forbidden = weights.masked_select(~allowed)
print(f"forbidden weight max: {forbidden.abs().max():.1e}")

i = 2
v2 = v.clone()
v2[:, :, i + 1:] = torch.randn_like(v2[:, :, i + 1:])
y1 = manual_attention(q, k, v, allowed)[0]
y2 = manual_attention(q, k, v2, allowed)[0]
print(f"max |y1-y2| over positions 0..{i}: {(y1[:, :, :i + 1] - y2[:, :, :i + 1]).abs().max():.2e}")
print(f"max |y1-y2| over positions {i + 1}..: {(y1[:, :, i + 1:] - y2[:, :, i + 1:]).abs().max():.4f}")

Expected outcome

good (2, 3, 5, 4)  bad (2, 3, 5, 4)  shapes equal: True
same dtype: True   same values: False
round trip, good: True   round trip, bad: False
good head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [12, 13, 14, 15],
        [24, 25, 26, 27],
        [36, 37, 38, 39],
        [48, 49, 50, 51]])
bad head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [16, 17, 18, 19]])
manual vs SDPA, max abs diff: 2.38e-07
row sums, max error from 1: 1.19e-07
forbidden weight max: 0.0e+00
max |y1-y2| over positions 0..2: 0.00e+00
max |y1-y2| over positions 3..: 2.4035

good (2, 3, 5, 4)  bad (2, 3, 5, 4)  shapes equal: True
same dtype: True   same values: False
round trip, good: True   round trip, bad: False
good head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [12, 13, 14, 15],
        [24, 25, 26, 27],
        [36, 37, 38, 39],
        [48, 49, 50, 51]])
bad head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [16, 17, 18, 19]])
manual vs SDPA, max abs diff: 2.38e-07
row sums, max error from 1: 1.19e-07
forbidden weight max: 0.0e+00
max |y1-y2| over positions 0..2: 0.00e+00
max |y1-y2| over positions 3..: 2.4035

Shape, then values, then behaviour

A shape assertion cannot see the bad split: it has the right shape. The chapter’s cheapest test is a round trip on a deterministic tensor: splitting and merging heads must return the original exactly.

def split_heads(x, num_heads):
    B, T, E = x.shape
    assert E % num_heads == 0, f"E={E} not divisible by num_heads={num_heads}"
    Dh = E // num_heads
    return x.reshape(B, T, num_heads, Dh).transpose(1, 2)

def split_heads(x, num_heads):
    B, T, E = x.shape
    assert E % num_heads == 0, f"E={E} not divisible by num_heads={num_heads}"
    Dh = E // num_heads
    return x.reshape(B, T, num_heads, Dh).transpose(1, 2)

The printed heads show the difference in arrangement. The good head 0 owns the first four embedding dimensions of every position. The bad head 0 owns the first twenty numbers of the flat sequence.

The last two lines are the behavioural test. Replacing the value vectors after i must not move the output at 0..i under a causal mask, and it does move the output after i. A mask that looks causal is a claim; an output that does not move is evidence.

Mechanism and limitations

Documented: the chapter shows two splits with identical shape, dtype and element multiset but different values, proves that merge_heads(split_heads(x)) is exact for the good split and not the bad one, and compares a manual implementation with PyTorch’s own.

Observed: the run above. Both splits had shape (2, 3, 5, 4) and the same dtype, but not the same values. The round trip was exact for the good split and false for the bad one. The manual attention matched PyTorch’s to 2.38e-07, the forbidden weight was 0.0e+00, and the intervention moved positions after i while leaving 0..i at zero.

Not shown. The program composes the chapter’s fragments into one run: it uses random Q, K and V instead of trained projections, and the intervention replaces value vectors rather than a trained model’s input. It checks one head count and one mask.

Canonical source , the chapter this example composes. Read the chapter .

Understand this example

Copy this prompt into your AI tool. No code runs here.


Canonical example context
Chapter: Attention: Which Position Is Comparing With Which?
http://programmer.ie/books/pytorch-from-first-principles/10-chapter/
Example: Two tensors, the same shape, and not the same tensor
http://programmer.ie/examples/pytorch-from-first-principles/10-chapter/


Attention is a sequence of explicit relationships between represented
vectors, and almost every stage can be wrong while producing exactly the
shape you predicted. Splitting the embedding across heads is the smallest
example: a reshape gives the shape you wanted and not the arrangement of
values you meant.

This program composes the chapter's pieces into one run: the good and bad
head split, a round trip that catches the bad one, a hand-written attention
compared against PyTorch's, and a future-token intervention.

## Run it

Save this as `attention_heads.py` and run it with Python 3.11 or later. It
needs PyTorch: `pip install torch`, then `python attention_heads.py`.

```python
"""Attention: which position is comparing with which, and how to prove it."""
import math

import torch
import torch.nn.functional as F

B, T, E, Nh = 2, 5, 12, 3
Dh = E // Nh


def split_heads(x, num_heads):
    B, T, E = x.shape
    return x.reshape(B, T, num_heads, E // num_heads).transpose(1, 2)


def merge_heads(x):
    B, Nh, T, Dh = x.shape
    return x.transpose(1, 2).reshape(B, T, Nh * Dh)


def manual_attention(q, k, v, allow):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    masked = scores.masked_fill(~allow, float("-inf"))
    weights = torch.softmax(masked, dim=-1)
    return weights @ v, weights


x = torch.arange(B * T * E).reshape(B, T, E)
good = split_heads(x, Nh)
bad = x.reshape(B, Nh, T, Dh)

print(f"good {tuple(good.shape)}  bad {tuple(bad.shape)}  shapes equal: {good.shape == bad.shape}")
print(f"same dtype: {good.dtype == bad.dtype}   same values: {torch.equal(good, bad)}")
print(f"round trip, good: {torch.equal(merge_heads(good), x)}   round trip, bad: {torch.equal(merge_heads(bad), x)}")
print("good head 0 of batch 0:")
print(good[0, 0])
print("bad head 0 of batch 0:")
print(bad[0, 0])

torch.manual_seed(0)
q, k, v = torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh)
allowed = torch.ones(T, T, dtype=torch.bool).tril()
mine, weights = manual_attention(q, k, v, allowed)
theirs = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
print(f"manual vs SDPA, max abs diff: {(mine - theirs).abs().max():.2e}")
print(f"row sums, max error from 1: {(weights.sum(-1) - 1).abs().max():.2e}")
forbidden = weights.masked_select(~allowed)
print(f"forbidden weight max: {forbidden.abs().max():.1e}")

i = 2
v2 = v.clone()
v2[:, :, i + 1:] = torch.randn_like(v2[:, :, i + 1:])
y1 = manual_attention(q, k, v, allowed)[0]
y2 = manual_attention(q, k, v2, allowed)[0]
print(f"max |y1-y2| over positions 0..{i}: {(y1[:, :, :i + 1] - y2[:, :, :i + 1]).abs().max():.2e}")
print(f"max |y1-y2| over positions {i + 1}..: {(y1[:, :, i + 1:] - y2[:, :, i + 1:]).abs().max():.4f}")
```

## Expected outcome

```text
good (2, 3, 5, 4)  bad (2, 3, 5, 4)  shapes equal: True
same dtype: True   same values: False
round trip, good: True   round trip, bad: False
good head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [12, 13, 14, 15],
        [24, 25, 26, 27],
        [36, 37, 38, 39],
        [48, 49, 50, 51]])
bad head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [16, 17, 18, 19]])
manual vs SDPA, max abs diff: 2.38e-07
row sums, max error from 1: 1.19e-07
forbidden weight max: 0.0e+00
max |y1-y2| over positions 0..2: 0.00e+00
max |y1-y2| over positions 3..: 2.4035
```

## Shape, then values, then behaviour

A shape assertion cannot see the bad split: it has the right shape. The
chapter's cheapest test is a round trip on a deterministic tensor: splitting
and merging heads must return the original exactly.

<!-- pytorch-from-first-principles-excerpt: content/books/pytorch-from-first-principles/10-chapter.md -->
```python
def split_heads(x, num_heads):
    B, T, E = x.shape
    assert E % num_heads == 0, f"E={E} not divisible by num_heads={num_heads}"
    Dh = E // num_heads
    return x.reshape(B, T, num_heads, Dh).transpose(1, 2)
```

The printed heads show the difference in arrangement. The good head 0 owns
the first four embedding dimensions of every position. The bad head 0 owns the
first twenty numbers of the flat sequence.

The last two lines are the behavioural test. Replacing the value vectors
after `i` must not move the output at `0..i` under a causal mask, and it does
move the output after `i`. A mask that looks causal is a claim; an output
that does not move is evidence.

## Mechanism and limitations

**Documented:** the chapter shows two splits with identical shape, dtype and
element multiset but different values, proves that
`merge_heads(split_heads(x))` is exact for the good split and not the bad
one, and compares a manual implementation with PyTorch's own.

**Observed:** the run above. Both splits had shape `(2, 3, 5, 4)` and the same
dtype, but not the same values. The round trip was exact for the good split
and false for the bad one. The manual attention matched PyTorch's to
`2.38e-07`, the forbidden weight was `0.0e+00`, and the intervention moved
positions after `i` while leaving `0..i` at zero.

**Not shown.** The program composes the chapter's fragments into one run: it
uses random Q, K and V instead of trained projections, and the intervention
replaces value vectors rather than a trained model's input. It checks one
head count and one mask.

[Canonical source](/books/pytorch-from-first-principles/10-chapter/), the
chapter this example composes.
[Read the chapter](https://programmer.ie/books/pytorch-from-first-principles/10-chapter/).
<!-- pytorch-from-first-principles-source: content/books/pytorch-from-first-principles/10-chapter.md sha256=6e48698ec89f9ff10d91f7eabe73e634a054037666c506a40cfc4cc3c33e7eb5 -->


Explain how this implementation demonstrates the chapter's mechanism, its expected outcome and its limitations.
Distinguish documented behaviour, observed behaviour and proposed choices. External AI may not be able to access these URLs, especially localhost or private previews. Use the source supplied here; ask for missing material rather than inventing it.

Apply this example

Copy this prompt into your AI tool. No code runs here.


Canonical example context
Chapter: Attention: Which Position Is Comparing With Which?
http://programmer.ie/books/pytorch-from-first-principles/10-chapter/
Example: Two tensors, the same shape, and not the same tensor
http://programmer.ie/examples/pytorch-from-first-principles/10-chapter/


Attention is a sequence of explicit relationships between represented
vectors, and almost every stage can be wrong while producing exactly the
shape you predicted. Splitting the embedding across heads is the smallest
example: a reshape gives the shape you wanted and not the arrangement of
values you meant.

This program composes the chapter's pieces into one run: the good and bad
head split, a round trip that catches the bad one, a hand-written attention
compared against PyTorch's, and a future-token intervention.

## Run it

Save this as `attention_heads.py` and run it with Python 3.11 or later. It
needs PyTorch: `pip install torch`, then `python attention_heads.py`.

```python
"""Attention: which position is comparing with which, and how to prove it."""
import math

import torch
import torch.nn.functional as F

B, T, E, Nh = 2, 5, 12, 3
Dh = E // Nh


def split_heads(x, num_heads):
    B, T, E = x.shape
    return x.reshape(B, T, num_heads, E // num_heads).transpose(1, 2)


def merge_heads(x):
    B, Nh, T, Dh = x.shape
    return x.transpose(1, 2).reshape(B, T, Nh * Dh)


def manual_attention(q, k, v, allow):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    masked = scores.masked_fill(~allow, float("-inf"))
    weights = torch.softmax(masked, dim=-1)
    return weights @ v, weights


x = torch.arange(B * T * E).reshape(B, T, E)
good = split_heads(x, Nh)
bad = x.reshape(B, Nh, T, Dh)

print(f"good {tuple(good.shape)}  bad {tuple(bad.shape)}  shapes equal: {good.shape == bad.shape}")
print(f"same dtype: {good.dtype == bad.dtype}   same values: {torch.equal(good, bad)}")
print(f"round trip, good: {torch.equal(merge_heads(good), x)}   round trip, bad: {torch.equal(merge_heads(bad), x)}")
print("good head 0 of batch 0:")
print(good[0, 0])
print("bad head 0 of batch 0:")
print(bad[0, 0])

torch.manual_seed(0)
q, k, v = torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh), torch.randn(B, Nh, T, Dh)
allowed = torch.ones(T, T, dtype=torch.bool).tril()
mine, weights = manual_attention(q, k, v, allowed)
theirs = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
print(f"manual vs SDPA, max abs diff: {(mine - theirs).abs().max():.2e}")
print(f"row sums, max error from 1: {(weights.sum(-1) - 1).abs().max():.2e}")
forbidden = weights.masked_select(~allowed)
print(f"forbidden weight max: {forbidden.abs().max():.1e}")

i = 2
v2 = v.clone()
v2[:, :, i + 1:] = torch.randn_like(v2[:, :, i + 1:])
y1 = manual_attention(q, k, v, allowed)[0]
y2 = manual_attention(q, k, v2, allowed)[0]
print(f"max |y1-y2| over positions 0..{i}: {(y1[:, :, :i + 1] - y2[:, :, :i + 1]).abs().max():.2e}")
print(f"max |y1-y2| over positions {i + 1}..: {(y1[:, :, i + 1:] - y2[:, :, i + 1:]).abs().max():.4f}")
```

## Expected outcome

```text
good (2, 3, 5, 4)  bad (2, 3, 5, 4)  shapes equal: True
same dtype: True   same values: False
round trip, good: True   round trip, bad: False
good head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [12, 13, 14, 15],
        [24, 25, 26, 27],
        [36, 37, 38, 39],
        [48, 49, 50, 51]])
bad head 0 of batch 0:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15],
        [16, 17, 18, 19]])
manual vs SDPA, max abs diff: 2.38e-07
row sums, max error from 1: 1.19e-07
forbidden weight max: 0.0e+00
max |y1-y2| over positions 0..2: 0.00e+00
max |y1-y2| over positions 3..: 2.4035
```

## Shape, then values, then behaviour

A shape assertion cannot see the bad split: it has the right shape. The
chapter's cheapest test is a round trip on a deterministic tensor: splitting
and merging heads must return the original exactly.

<!-- pytorch-from-first-principles-excerpt: content/books/pytorch-from-first-principles/10-chapter.md -->
```python
def split_heads(x, num_heads):
    B, T, E = x.shape
    assert E % num_heads == 0, f"E={E} not divisible by num_heads={num_heads}"
    Dh = E // num_heads
    return x.reshape(B, T, num_heads, Dh).transpose(1, 2)
```

The printed heads show the difference in arrangement. The good head 0 owns
the first four embedding dimensions of every position. The bad head 0 owns the
first twenty numbers of the flat sequence.

The last two lines are the behavioural test. Replacing the value vectors
after `i` must not move the output at `0..i` under a causal mask, and it does
move the output after `i`. A mask that looks causal is a claim; an output
that does not move is evidence.

## Mechanism and limitations

**Documented:** the chapter shows two splits with identical shape, dtype and
element multiset but different values, proves that
`merge_heads(split_heads(x))` is exact for the good split and not the bad
one, and compares a manual implementation with PyTorch's own.

**Observed:** the run above. Both splits had shape `(2, 3, 5, 4)` and the same
dtype, but not the same values. The round trip was exact for the good split
and false for the bad one. The manual attention matched PyTorch's to
`2.38e-07`, the forbidden weight was `0.0e+00`, and the intervention moved
positions after `i` while leaving `0..i` at zero.

**Not shown.** The program composes the chapter's fragments into one run: it
uses random Q, K and V instead of trained projections, and the intervention
replaces value vectors rather than a trained model's input. It checks one
head count and one mask.

[Canonical source](/books/pytorch-from-first-principles/10-chapter/), the
chapter this example composes.
[Read the chapter](https://programmer.ie/books/pytorch-from-first-principles/10-chapter/).
<!-- pytorch-from-first-principles-source: content/books/pytorch-from-first-principles/10-chapter.md sha256=6e48698ec89f9ff10d91f7eabe73e634a054037666c506a40cfc4cc3c33e7eb5 -->


Help me adapt this example to my project. First identify the mechanism. Ask only for project details necessary to assess applicability. Prefer the smallest useful implementation. Explain adaptations before making changes and define verification. The original example does not prove an adaptation works.
Distinguish documented behaviour, observed behaviour and proposed choices. External AI may not be able to access these URLs, especially localhost or private previews. Use the source supplied here; ask for missing material rather than inventing it.