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

import mlx.core as mx
from dataclasses import dataclass
from mlx import nn

@dataclass
class LayerAttentionData:
    layer: int
    img_attention: mx.array
    concept_attention: mx.array
    ...

class JointTransformerBlockConcept(nn.Module):
    def __init__(self, layer) -> None: ...
    def __call__(
        self,
        layer_idx: int,
        hidden_states: mx.array,
        encoder_hidden_states: mx.array,
        encoder_hidden_states_concept: mx.array,
        text_embeddings: mx.array,
        text_embeddings_concept: mx.array,
        rotary_embeddings: mx.array,
        rotary_embeddings_concept: mx.array,
    ) -> tuple[mx.array, mx.array, mx.array, LayerAttentionData]: ...
