N/C ？

[Mrfanl]

if (N / C) < 1: x = ((q @ k.transpose(-2, -1)) @ v).transpose(1, 2).reshape(B, N, C) else: x = (q @ (k.transpose(-2, -1) @ v)).transpose(1, 2).reshape(B, N, C)
should be
if (N / （C/H）) < 1: x = ((q @ k.transpose(-2, -1)) @ v).transpose(1, 2).reshape(B, N, C) else: x = (q @ (k.transpose(-2, -1) @ v)).transpose(1, 2).reshape(B, N, C)
？

[soroush-abbasi]

Hi,

Sorry for the late reply. You are correct. Thanks for pointing out this bug!

BTW, in the paper we show that SimA works well even without multi-head attention (single head H=1).