Attention correction

Fixing attention. Added a basic test value.
The output `weight` does not have the same axes as the input  `mask`. Is this desirable ?