Add a batch_first argument to MultiheadAttention

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 moved

Checking 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.

Start at the pytensor_ml.layers.MultiheadAttention entry point and reproduce the supplied NumPy example with the sequence-first shape. Trace how the embedding axis is validated, then add coverage showing that an incompatible layout is rejected; done means the demonstrated cross-batch mixing cannot occur silently.