`.typings/mflux/models/qwen/model/qwen_transformer/qwen_attention.pyi`

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

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

class QwenAttention(nn.Module):
    _num_heads: int
    _head_dim: int
    num_heads: int
    head_dim: int
    to_q: nn.Linear
    to_k: nn.Linear
    to_v: nn.Linear
    add_q_proj: nn.Linear
    add_k_proj: nn.Linear
    add_v_proj: nn.Linear
    norm_q: nn.RMSNorm
    norm_k: nn.RMSNorm
    norm_added_q: nn.RMSNorm
    norm_added_k: nn.RMSNorm
    attn_to_out: list[Any]
    to_add_out: nn.Linear

    def __init__(
        self, dim: int = ..., num_heads: int = ..., head_dim: int = ...
    ) -> None: ...
    def __call__(
        self,
        img_modulated: mx.array,
        txt_modulated: mx.array,
        encoder_hidden_states_mask: mx.array | None,
        image_rotary_emb: tuple[mx.array, mx.array],
        block_idx: int | None = ...,
    ) -> tuple[mx.array, mx.array]: ...
    def _compute_attention_qwen(
        self,
        query: mx.array,
        key: mx.array,
        value: mx.array,
        mask: mx.array | None,
        block_idx: int | None,
    ) -> mx.array: ...
    @staticmethod
    def _convert_mask_for_qwen(
        mask: mx.array | None, joint_seq_len: int, txt_seq_len: int
    ) -> mx.array | None: ...
    @staticmethod
    def _apply_rope_qwen(
        x: mx.array, cos_vals: mx.array, sin_vals: mx.array
    ) -> mx.array: ...