efficient implementation of attention

[Coerulatus]

Hello,
@levtelyatnikov and I made some changes to allset_transformer_layer.
We noticed that the old implementation required a lot of memory when calculating the attention scores when dealing with large graphs. Now the layer achieves the same results with less memory requirements.
We also introduced a LeakyReLU activation when calculating the attention as per the official code implementation (lines 168-177).
Thank you very much.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.37%. Comparing base (2267768) to head (d702faf).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #275   +/-   ##
=======================================
  Coverage   97.37%   97.37%           
=======================================
  Files          58       58           
  Lines        2054     2058    +4     
=======================================
+ Hits         2000     2004    +4     
  Misses         54       54           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ninamiolane]

Great catch, thank you!