Function from_head_matmul 

pub fn from_head_matmul(
    spec: &HeadMatmulSpec,
    names: &StageNames,
    target: &SpatialTarget,
) -> Result<SpatialProgram, SpatialReject>

Y = X matmul W across H heads — the weighted-sum shape: the folded axis is the ROW (e over a head’s a inputs), the flat output column c = u·b + r for unit u (row u / H, head h = u mod H), and each unit’s band reads X[u·a + e] and W[(h·a + e)·b + c − u·b].