def create_self_attention(num_channels: int, num_heads: int, dropout: float, activation_checkpoint: bool = False):
    """
    Self attention block for Perceiver https://arxiv.org/pdf/2103.03206.pdf.

    Args:
        num_channels (`int`):
            Number of q channels.
        num_heads (`int`):
            Number of parallel attention heads.
        dropout (`nn.Module`):
            Dropout probability.
        activation_checkpoint (`bool`):
            Use activation checkpointing.

    Return (`nn.Module`):
        Configured self attention layer.
    """
    layer = Sequential(
        Residual(SelfAttention(num_channels, num_heads, dropout), dropout), Residual(mlp(num_channels), dropout)
    )
    return layer if not activation_checkpoint else checkpoint_wrapper(layer)
Self attention block for Perceiver https://arxiv.org/pdf/2103.03206.pdf.

Args:
num_channels (`int`):
Number of q channels.
num_heads (`int`):
Number of parallel attention heads.
dropout (`nn.Module`):
Dropout probability.
activation_checkpoint (`bool`):
Use activation checkpointing.

Return (`nn.Module`):
Configured self attention layer.
towhee/models/perceiver/create_self_attention.py