Einops

Mental Model

# Old Mental Model (Karpathy)
# 'B T C -> ...'

# New Mental Model (ARENA)
# 'b s d_model -> ...'

Rearrange

Core Syntax

output = rearrange(tensor, 'input_pattern -> output_pattern', **constants)

Move 1: The Swap (Permute)

• What it does: Reorders dimensions (like .transpose or .permute).
• When to use it: Moving the heads dimension next to the batch dimension so you can parallelize attention.
• Visual: You are just rotating the cube.

# Move 'c' (channels) to the front
y = rearrange(x, 'b s c -> c b s')

Move 2: The Split (Decomposition)

• What it does: Breaks one dimension into two (or more).When to use it: Turning d_model (C) into n_heads and d_head