def forward(self, src, attn_mask=None, key_padding_mask=None, state=None, **kwargs):
        """ state should represent a mask and key padding mask """
        if self.causal and attn_mask is None:
            attn_mask = TriangularCausalMask(src.size(-2), device=src.device)
        # attn_mask, key_padding_mask = state
        # Note that this returns None for the second argument
        y, z = self.mha(src, src, src, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=False)
        return y, None
 state should represent a mask and key padding mask 
analyze/convnexts4nd/src/models/sequence/attention/linear.py