Repository navigation
Add a batch_first argument to MultiheadAttention #127
Description
torch.nn.MultiheadAttention defaults to batch_first=False, so sequence-first is the layout a torch user arrives holding. Here the last axis is the embedding and the second-to-last is treated as the sequence, so a seq-first batch is legal, produces an output of the same shape, and mixes every example in the batch into every other one. With is_causal=True you additionally get a causal mask over the batch axis. Nothing raises at any point.
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor_ml.layers import MultiheadAttention
mha = MultiheadAttention("a", n_embd=8, n_head=2)
X = pt.tensor("X", shape=(5, 3, 8)) # torch's default: (seq, batch, embd)
f = pytensor.function([X], mha(X))
rng = np.random.default_rng(0)
a = rng.normal(size=(5, 3, 8))
b = a.copy()
b[:, 1, :] += 10.0 # perturb batch element 1 only
print(np.abs(f(a)[:, 0, :] - f(b)[:, 0, :]).max()) # 0.708 -- element 0 movedChecking that the embedding axis matches n_embd would catch it, since the seq-first layout only type-checks when n_embd happens to equal the batch size.
Activity
[-]MultiheadAttention on a (seq, batch, embd) input attends across the batch[/-][+]Add a batch_first argument to MultiheadAttention[/+]jessegrabowski commentedon Sep 1, 2026
The layer documents
(batch, time, n_embd), and both layouts are legal shapes, so nothing can tell a seq-first input from a batch-first one. That makes this a missing feature rather than a bug. What's worth doing is abatch_firstargument, the way torch spells it.