Nós não fazemos qualquer IA.Fazemos a sua IA.Do seu jeito.

"""Multi-head self-attention — the core operation behind modern LLMs."""

from __future__ import annotations

import torch
from torch import Tensor, nn
from torch.nn import functional as F

__all__ = ["MultiHeadSelfAttention"]


class MultiHeadSelfAttention(nn.Module):
    """Multi-head self-attention over a sequence of token embeddings.

    Computes scaled dot-product attention (Vaswani et al., 2017) through
    PyTorch's fused scaled_dot_product_attention kernel, which dispatches
    to FlashAttention / memory-efficient backends with O(n) memory.

    Args:
        embed_dim: Total embedding dimension, split evenly across heads.
        num_heads: Number of attention heads; must divide embed_dim.
        dropout: Dropout probability applied to the attention weights
            during training.
    """

    def __init__(
        self, embed_dim: int, num_heads: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        if embed_dim % num_heads != 0:
            raise ValueError(
                f"embed_dim ({embed_dim}) must be divisible by "
                f"num_heads ({num_heads})."
            )

        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.dropout_p = dropout

        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self._reset_parameters()

    def _reset_parameters(self) -> None:
        # Xavier/Glorot init for attention projections (Glorot & Bengio,
        # 2010), as in nn.MultiheadAttention; Linear's default Kaiming init
        # targets ReLU fan-in, not softmax attention.
        for projection in (self.qkv_proj, self.out_proj):
            nn.init.xavier_uniform_(projection.weight)
            nn.init.zeros_(projection.bias)

    def forward(
        self,
        x: Tensor,
        mask: Tensor | None = None,
        is_causal: bool = False,
    ) -> Tensor:
        """Apply self-attention.

        Args:
            x: Embeddings of shape (batch, seq_len, embed_dim).
            mask: Optional boolean mask broadcastable to
                (batch, num_heads, seq_len, seq_len); True lets a
                position attend, False blocks it.
            is_causal: If True, apply a causal (lower-triangular) mask so
                each position attends only to itself and earlier positions.

        Returns:
            Tensor of shape (batch, seq_len, embed_dim).
        """
        batch, seq_len, embed_dim = x.shape

        # One projection, then split into query, key and value per head.
        qkv = self.qkv_proj(x)
        qkv = qkv.reshape(batch, seq_len, 3, self.num_heads, self.head_dim)
        query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(dim=0)

        # Fused scaled dot-product attention (FlashAttention backend).
        # Scaling by 1/sqrt(head_dim) happens inside the kernel; dropout
        # must be disabled outside training.
        context = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=mask,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=is_causal,
        )

        # Concatenate heads and project back to the model dimension.
        context = context.permute(0, 2, 1, 3).reshape(batch, seq_len, embed_dim)
        return self.out_proj(context)