TypeError: forward() missing 3 required positional arguments: 'attn_mask', 'query_lengths', and 'key_lengths'

[TianhaoFu]

Hi, In fast-transformers/fast_transformers/attention/clustered_attention.py , I found that clusted attention forward function like this :

 def forward(self, queries, keys, values, attn_mask, query_lengths,
                key_lengths):

But in fast-transformers/fast_transformers/attention/attention_layer.py

 def test_full_attention_forward(self):
        d_model = 128
        n_heads = 4
        transformer = TransformerEncoder([
            TransformerEncoderLayer(
                AttentionLayer(
                    ClusteredAttention(
                        clusters = 10
                    ),
                    d_model,
                    n_heads
                ),
                d_model,
                n_heads
            )
            for i in range(6)
        ])

        transformer = transformer.to("cuda")
        x = torch.rand(100, 20, d_model).to("cuda")
        y = transformer(x)

It only pass a parameter x and don't pass 'attn_mask', 'query_lengths', and 'key_lengths'.

Is this a bug? Thanks

[apoorv2904]

Hi,

This is not a bug. In this case the attention module is called by TransformerEncoder which accepts attn_mask and length_mask as optional inputs and appropriately calls the underlying attention. Please find the documentation here and for more details please take look at the forward pass of TransformerEncoder.

I will close the issue. Feel free to open for more questions.