# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/layers/01_attention.ipynb.

# %% auto 0
__all__ = ['Attention1DConfig', 'Attention3DConfig', 'Attention1D', 'Attention3D']

# %% ../../nbs/layers/01_attention.ipynb 2
from functools import partial, wraps

import torch
import torch.nn.functional as F
from einops import rearrange
from torch import nn

from ..docstrings import populate_docstring
from .embeddings import RelativePositionEmbeddings
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import CustomBaseModel, Field, model_validator
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/layers/01_attention.ipynb 4
class Attention1DConfig(CustomBaseModel):
    dim: int | tuple[int, int] = Field(
        ...,
        description=(
            "Dimension of the input features. If tuple, (dim_qk, dim_v). "
            "Otherwise it is assumed to be dim of both qk and v."
        ),
    )
    num_heads: int = Field(..., description="Number of query heads")
    ratio_q_to_kv_heads: int = Field(1, description="Ratio of query heads to key/value heads. Useful for MQA/GQA.")
    logit_scale_learnable: bool = Field(False, description="Whether the logit scale is learnable.")
    attn_drop_prob: float = Field(0.0, description="Dropout probability for attention weights.")
    proj_drop_prob: float = Field(0.0, description="Dropout probability for the projection layer.")
    max_attention_batch_size: int = Field(
        -1,
        description=(
            "Runs attention by splitting the inputs into chunks of this size. 0 means no chunking. "
            "Useful for large inputs during inference. (This happens along batch dimension)."
        ),
    )

    @property
    def num_q_heads(self) -> int:
        return self.num_heads

    @property
    def num_kv_heads(self) -> int:
        return self.num_heads // self.ratio_q_to_kv_heads

    @property
    def gqa_mqa_enabled(self) -> bool:
        return self.ratio_q_to_kv_heads != 1

    @property
    def dim_qk(self) -> int:
        if isinstance(self.dim, tuple):
            return self.dim[0]
        return self.dim

    @property
    def dim_v(self) -> int:
        if isinstance(self.dim, tuple):
            return self.dim[1]
        return self.dim

    @property
    def per_head_dim_qk(self) -> int:
        return self.dim_qk // self.num_heads

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        if self.gqa_mqa_enabled:
            assert torch.__version__ >= "2.5", "Need PyTorch version >= 2.5 for GQA and MQA"

        assert self.dim_qk % self.num_heads == 0, "dimension must be divisible by number of heads"
        assert (
            self.num_heads % self.num_kv_heads == 0
        ), "number of query heads must be divisible by number of key and value heads"

        return self


class Attention3DConfig(Attention1DConfig):
    pass

# %% ../../nbs/layers/01_attention.ipynb 6
@populate_docstring
class Attention1D(nn.Module):
    """Performs attention (MHA, GQA, and MQA) on 1D sequences. {CLASS_DESCRIPTION_1D_DOC}"""

    _warn_relative_position_bias: bool = True

    @populate_docstring
    def __init__(
        self,
        config: Attention1DConfig = {},
        relative_position_bias: RelativePositionEmbeddings | None = None,
        logit_scale: float | None = None,
        checkpointing_level: int = 0,
        **kwargs
    ):
        """Initializes the Attention1D module.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            relative_position_bias: Relative position embeddings to be considered during attention. Should be callable.
            logit_scale: Logit scale to be used for attention. If None, it will be initialized based on per-head
                dimension.
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = Attention1DConfig.model_validate(config | kwargs)

        dim_qk = self.config.dim_qk
        dim_v = self.config.dim_v
        ratio_q_to_kv_heads = self.config.ratio_q_to_kv_heads
        per_head_dim = self.config.per_head_dim_qk
        logit_scale_learnable = self.config.logit_scale_learnable
        attn_drop_prob = self.config.attn_drop_prob
        proj_drop_prob = self.config.proj_drop_prob

        self.W_q = nn.Linear(dim_qk, dim_qk)
        self.W_k = nn.Linear(dim_qk, dim_qk // ratio_q_to_kv_heads)
        self.W_v = nn.Linear(dim_v, dim_qk // ratio_q_to_kv_heads)
        self.attn_drop_prob = attn_drop_prob
        self.proj = nn.Linear(dim_qk, dim_qk)
        self.proj_drop = nn.Dropout(proj_drop_prob)

        if logit_scale is None:
            self.logit_scale = nn.Parameter(
                torch.tensor([per_head_dim**-0.5]),
                requires_grad=logit_scale_learnable,
            )
        else:
            self.logit_scale = logit_scale

        if self._warn_relative_position_bias and relative_position_bias is not None:
            print(
                "Warning: Relative position bias is not used in Attention1D. "
                "Use Attention3D for relative position bias."
            )
        self.relative_position_bias = relative_position_bias

        self.checkpointing_level1 = ActivationCheckpointing(1, checkpointing_level)
        self.checkpointing_level2 = ActivationCheckpointing(2, checkpointing_level)

    def _forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
        """Forward pass of the Attention1D module.

        Terminology: T => number of tokens, b => batch size

        Args:
            query: Tensor of shape (b, T_q, dim_qk) representing the input to the query matrix.
            key: Tensor of shape (b, T_kv, dim_qk) representing the input to the key matrix.
            value: Tensor of shape (b, T_kv, dim_v) representing the input to the value matrix.

        Returns:
            Tensor of shape (b, T_q, dim_qk) representing output tokens.
        """

        def get_final_query_key_value(query, key, value):
            """Computing query, key, and value tokens after passing to the weight matrices. Useful for activation
            checkpointing"""
            query = self.W_q(query)
            key = self.W_k(key)
            value = self.W_v(value)

            rearrange_partial = partial(rearrange, pattern="b T (num_heads d) -> b num_heads T d")
            query = rearrange_partial(query, num_heads=self.config.num_heads).contiguous()
            key = rearrange_partial(key, num_heads=self.config.num_kv_heads).contiguous()
            value = rearrange_partial(value, num_heads=self.config.num_kv_heads).contiguous()
            # query: (b, num_heads, T, per_head_dim)
            # key: (b, num_kv_heads, T, per_head_dim)
            # value: (b, num_kv_heads, T, per_head_dim)

            if isinstance(self.logit_scale, nn.Module):
                logit_scale = self.logit_scale()
            else:
                logit_scale = self.logit_scale

            query_normalized = F.normalize(query, dim=-1)
            key_normalized = F.normalize(key, dim=-1)

            query_normalized_and_scaled = query_normalized * logit_scale  # Scale the query beforehand

            return query_normalized_and_scaled, key_normalized, value

        query_normalized_and_scaled, key_normalized, value = self.checkpointing_level1(
            get_final_query_key_value, query, key, value
        )

        relative_position_bias = None
        if self.relative_position_bias is not None:
            relative_position_bias = self.relative_position_bias()

        # Split tensors into batches and perform attention
        output = []
        chunk_size = self.config.max_attention_batch_size
        if chunk_size == -1:
            chunk_size = query_normalized_and_scaled.size(0)
        for query_normalized_and_scaled_chunk, key_normalized_chunk, value_chunk in zip(
            torch.split(query_normalized_and_scaled, chunk_size, dim=0),
            torch.split(key_normalized, chunk_size, dim=0),
            torch.split(value, chunk_size, dim=0),
        ):
            torch250plus_kwargs = {}
            if torch.__version__ >= "2.5":
                torch250plus_kwargs["enable_gqa"] = self.config.gqa_mqa_enabled

            output_chunk = F.scaled_dot_product_attention(
                query_normalized_and_scaled_chunk,
                key_normalized_chunk,
                value_chunk,
                attn_mask=relative_position_bias,  # Use this as a way to introduce relative position bias
                dropout_p=self.attn_drop_prob,
                is_causal=False,
                scale=1.0,  # Already scaled the vectors
                **torch250plus_kwargs,
            )
            output.append(output_chunk)
            # (chunk_size, num_heads, T, per_head_dim)
        output = torch.cat(output, dim=0)
        # (b, num_heads, T, per_head_dim)

        output = rearrange(output, "b num_heads T d -> b T (num_heads d)").contiguous()
        # (b, T, dim_qk)

        def get_final_output(output):
            """Computing final output after projection. Useful for activation checkpointing"""
            output = self.proj(output)
            output = self.proj_drop(output)
            return output

        output = self.checkpointing_level1(get_final_output, output)
        # (b, T, dim_qk)

        return output

    @wraps(_forward)
    def forward(self, *args, **kwargs):
        return self.checkpointing_level2(self._forward, *args, **kwargs)

# %% ../../nbs/layers/01_attention.ipynb 8
@populate_docstring
class Attention3D(Attention1D):
    """Performs attention (MHA, GQA, and MQA) on 3D sequences. {CLASS_DESCRIPTION_3D_DOC}"""

    _warn_relative_position_bias: bool = False

    def __init__(
        self,
        config: Attention3DConfig = {},
        relative_position_bias: RelativePositionEmbeddings | None = None,
        logit_scale: float | None = None,
        checkpointing_level: int = 0,
        **kwargs
    ):
        super().__init__(config, relative_position_bias, logit_scale, checkpointing_level, **kwargs)

    @populate_docstring
    def _forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        channels_first: bool = True,
    ):
        """Forward pass of the Attention3D module.

        Terminology: z => depth, y => height, x => width, b => batch size

        Args:
            query: Tensor of shape (b, [dim_qk], z_q, y_q, x_q, [dim_qk]) representing the input to the query matrix.
            key: Tensor of shape (b, [dim_qk], z_kv, y_kv, x_kv, [dim_qk]) representing the input to the key matrix.
            value: Tensor of shape (b, [dim_v], z_kv, y_kv, x_kv, [dim_v]) representing the input to the value matrix.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            Tensor of shape (b, [dim_qk], z_q, y_q, x_q, [dim_qk]) representing output tokens.

        Constraints:
            z_q * y_q * x_q = z_kv * y_kv * x_kv
        """
        query = rearrange_channels(query, channels_first, False)
        key = rearrange_channels(key, channels_first, False)
        value = rearrange_channels(value, channels_first, False)
        # Each is now (b, z, y, x, d)

        z_q, y_q, x_q = query.shape[1:4]

        query = rearrange(query, "b z y x d -> b (z y x) d").contiguous()
        key = rearrange(key, "b z y x d -> b (z y x) d").contiguous()
        value = rearrange(value, "b z y x d -> b (z y x) d").contiguous()
        # Each is now (b, T, d)

        output = super()._forward(query, key, value)
        # (b, T, d)

        output = rearrange(output, "b (z y x) d -> b z y x d", z=z_q, y=y_q, x=x_q).contiguous()
        # (b, z, y, x, d)

        output = rearrange_channels(output, False, channels_first)
        # (b, [d], z, y, x, [d])

        return output

    @wraps(_forward)
    def forward(self, *args, **kwargs):
        return self.checkpointing_level2(self._forward, *args, **kwargs)
