def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = True,
        use_rel_pos: bool = False,
        rel_pos_zero_init: bool = True,
        input_size: Optional[tuple[int, int]] = None,
    ) -> None:
        """Construct attention block.

        Args:
        dim: Number of input channels.
        num_heads: Number of attention heads.
        qkv_bias:  If True, add a learnable bias to query, key, value.
        use_rel_pos: If True, add relative positional embeddings to the attention map.
        rel_pos_zero_init: If True, zero initialize relative positional parameters.
        input_size: Input resolution for calculating the relative positional parameter size.

        """
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

        self.use_rel_pos = use_rel_pos
        if self.use_rel_pos and isinstance(input_size, tuple):
            # initialize relative positional embeddings
            self.rel_pos_h = nn.Parameter(zeros(2 * input_size[0] - 1, head_dim))
            self.rel_pos_w = nn.Parameter(zeros(2 * input_size[1] - 1, head_dim))
Construct attention block.

        Args:
        dim: Number of input channels.
        num_heads: Number of attention heads.
        qkv_bias:  If True, add a learnable bias to query, key, value.
        use_rel_pos: If True, add relative positional embeddings to the attention map.
        rel_pos_zero_init: If True, zero initialize relative positional parameters.
        input_size: Input resolution for calculating the relative positional parameter size.

        
kornia/contrib/models/sam/architecture/image_encoder.py