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

import mlx.core as mx
from mlx import nn
from typing import Any

class JointAttention(nn.Module):
    num_heads: int
    head_dimension: int
    to_q: nn.Linear
    to_k: nn.Linear
    to_v: nn.Linear
    norm_q: nn.RMSNorm
    norm_k: nn.RMSNorm
    add_q_proj: nn.Linear
    add_k_proj: nn.Linear
    add_v_proj: nn.Linear
    norm_added_q: nn.RMSNorm
    norm_added_k: nn.RMSNorm
    to_out: list[Any]
    to_add_out: nn.Linear

    def __init__(self) -> None: ...
    def __call__(
        self,
        hidden_states: mx.array,
        encoder_hidden_states: mx.array,
        image_rotary_emb: mx.array,
    ) -> tuple[mx.array, mx.array]: ...
