Needed for #34 . In multi head attention, you concatenate the output of all the single attention heads.

Note that this may require nightly access similar to #1 because we can't do expressions with const generics yet.