"""Lesson 6: multi-head attention and the transformer block - check the shapes and the causal mask."""
import torch
from tinygpt import MultiHeadAttention, Block
torch.manual_seed(0)
x = torch.randn(2, 8, 32)                      # batch 2, 8 tokens, 32-dim embeddings
mha = MultiHeadAttention(d_model=32, n_heads=4, block_size=16)
print("attention out:", tuple(mha(x).shape))  # same shape as input
block = Block(d_model=32, n_heads=4, block_size=16)
print("block out:", tuple(block(x).shape))
# Causal check: changing a FUTURE token must not change earlier outputs.
y1 = block(x); x2 = x.clone(); x2[:, -1] += 10; y2 = block(x2)
print("earlier positions unchanged:", torch.allclose(y1[:, :-1], y2[:, :-1], atol=1e-6))
print("parameters in one block:", sum(p.numel() for p in block.parameters()))