Merge branch 'develop' into feature/44-make-flash-attention-configurable

src/anemoi/models/layers/attention.py

@@ -34,6 +34,7 @@ def __init__(
         window_size: Optional[int] = None,
         dropout_p: float = 0.0,
         softcap: float = 0.0,
+
     ):
         super().__init__()
 

tests/layers/block/test_block_transformer.py

@@ -38,6 +38,7 @@ def test_init(self, factor_attention_heads, hidden_dim, num_heads, activation, w
         num_channels = num_heads * factor_attention_heads
         block = TransformerProcessorBlock(
             num_channels, hidden_dim, num_heads, activation, window_size, dropout_p=dropout_p, softcap=softcap
+
         )
         assert isinstance(block, TransformerProcessorBlock)
 

tests/layers/test_attention.py

@@ -48,6 +48,7 @@ def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_mult
     embed_dim = num_heads * embed_dim_multiplier
     mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p, softcap=softcap)
 
+
     x = torch.randn(batch_size * 2, embed_dim)
     shapes = [list(x.shape)]
     output = mhsa.forward(x, shapes, batch_size)
@@ -68,6 +69,7 @@ def test_multi_head_self_attention_backward(batch_size, num_heads, embed_dim_mul
     embed_dim = num_heads * embed_dim_multiplier
     mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p, softcap=softcap)
 
+
     x = torch.randn(batch_size * 2, embed_dim, requires_grad=True)
     shapes = [list(x.shape)]
     output = mhsa.forward(x, shapes, batch_size)