"""
This type stub file was generated by pyright.
"""

from typing import Any, Callable, Optional

import mlx.core as mx
from base import Module

class MultiHeadAttention(Module):
    """Implements the scaled dot product attention with multiple heads.

    Given inputs for queries, keys and values the ``MultiHeadAttention``
    produces new values by aggregating information from the input values
    according to the similarities of the input queries and keys.

    All inputs as well as the output are linearly projected without biases by
    default.

    ``MultiHeadAttention`` also takes an optional additive attention mask that
    should be broadcastable with ``(batch, num_heads, # queries, # keys)``. The
    mask should have ``-inf`` or very large negative numbers at the positions
    that should *not* be attended to.

    Args:
        dims (int): The model dimensions. This is also the default
            value for the queries, keys, values, and the output.
        num_heads (int): The number of attention heads to use.
        query_input_dims (int, optional): The input dimensions of the queries.
            Default: ``dims``.
        key_input_dims (int, optional): The input dimensions of the keys.
            Default: ``dims``.
        value_input_dims (int, optional): The input dimensions of the values.
            Default: ``key_input_dims``.
        value_dims (int, optional): The dimensions of the values after the
            projection. Default: ``dims``.
        value_output_dims (int, optional): The dimensions the new values will
            be projected to. Default: ``dims``.
        bias (bool, optional): Whether or not to use a bias in the projections.
            Default: ``False``.
    """
    def __init__(
        self,
        dims: int,
        num_heads: int,
        query_input_dims: Optional[int] = ...,
        key_input_dims: Optional[int] = ...,
        value_input_dims: Optional[int] = ...,
        value_dims: Optional[int] = ...,
        value_output_dims: Optional[int] = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(
        self, queries: mx.array, keys: mx.array, values: mx.array, mask: mx.array = ...
    ) -> mx.array: ...
    @staticmethod
    def create_additive_causal_mask(N: int, dtype: mx.Dtype = ...) -> mx.array: ...

class TransformerEncoderLayer(Module):
    def __init__(
        self,
        dims: int,
        num_heads: int,
        mlp_dims: Optional[int] = ...,
        dropout: float = ...,
        activation: Callable[[Any], Any] = ...,
        norm_first: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ...

class TransformerEncoder(Module):
    def __init__(
        self,
        num_layers: int,
        dims: int,
        num_heads: int,
        mlp_dims: Optional[int] = ...,
        dropout: float = ...,
        activation=...,
        norm_first: bool = ...,
        checkpoint: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ...

class TransformerDecoderLayer(Module):
    def __init__(
        self,
        dims: int,
        num_heads: int,
        mlp_dims: Optional[int] = ...,
        dropout: float = ...,
        activation: Callable[[Any], Any] = ...,
        norm_first: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ...

class TransformerDecoder(Module):
    def __init__(
        self,
        num_layers: int,
        dims: int,
        num_heads: int,
        mlp_dims: Optional[int] = ...,
        dropout: float = ...,
        activation=...,
        norm_first: bool = ...,
        checkpoint: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ...

class Transformer(Module):
    """
    Implements a standard Transformer model.

    The implementation is based on `Attention Is All You Need
    <https://arxiv.org/abs/1706.03762>`_.

    The Transformer model contains an encoder and a decoder. The encoder
    processes the input sequence and the decoder generates the output sequence.
    The interaction between encoder and decoder happens through the attention
    mechanism.

    Args:
        dims (int, optional): The number of expected features in the
            encoder/decoder inputs. Default: ``512``.
        num_heads (int, optional): The number of attention heads. Default:
            ``8``.
        num_encoder_layers (int, optional): The number of encoder layers in the
            Transformer encoder. Default: ``6``.
        num_decoder_layers (int, optional): The number of decoder layers in the
            Transformer decoder. Default: ``6``.
        mlp_dims (int, optional): The hidden dimension of the MLP block in each
            Transformer layer. Defaults to ``4*dims`` if not provided. Default:
            ``None``.
        dropout (float, optional): The dropout value for the Transformer
            encoder and decoder. Dropout is used after each attention layer and
            the activation in the MLP layer. Default: ``0.0``.
        activation (function, optional): the activation function for the MLP
            hidden layer. Default: :func:`relu`.
        custom_encoder (nn.Module, optional): A custom encoder to replace the
            standard Transformer encoder. Default: ``None``.
        custom_decoder (nn.Module, optional): A custom decoder to replace the
            standard Transformer decoder. Default: ``None``.
        norm_first (bool, optional): if ``True``, encoder and decoder layers
            will perform layer normalization before attention and MLP
            operations, otherwise after. Default: ``True``.
        checkpoint (bool, optional): if ``True`` perform gradient checkpointing
            to reduce the memory usage at the expense of more computation.
            Default: ``False``.
    """
    def __init__(
        self,
        dims: int = ...,
        num_heads: int = ...,
        num_encoder_layers: int = ...,
        num_decoder_layers: int = ...,
        mlp_dims: Optional[int] = ...,
        dropout: float = ...,
        activation: Callable[[Any], Any] = ...,
        custom_encoder: Optional[Any] = ...,
        custom_decoder: Optional[Any] = ...,
        norm_first: bool = ...,
        checkpoint: bool = ...,
    ) -> None: ...
    def __call__(
        self, src, tgt, src_mask, tgt_mask, memory_mask
    ) -> mx.array:  # -> array | Any:
        ...
