def __call__(
        self,
        tgt_ids_BxTxC: mx.array,
        encoder_out: mx.array,
        tgt_positions: mx.array,
        src_positions: mx.array,
        deterministic: bool,
        self_attn_mask: mx.array,
        cross_attn_mask: mx.array,
        self_attention_cache: List[KVCache],
        cross_attention_cache: List[KVCache],
    ) -> mx.array:
        """
        Forward pass for the Decoder stack, managing KV caches.

        Args:
            tgt_ids_BxTxC: Target token IDs (B, T, C).
            encoder_out: Output from the encoder (B, S, E).
            tgt_positions: Positions for target sequence (B, T).
            src_positions: Positions for source sequence (B, S).
            deterministic: Disable dropout if True.
            self_attn_mask: Mask for self-attention.
            cross_attn_mask: Mask for cross-attention.
            self_attention_cache: List containing the self-attention KV cache for each layer.
            cross_attention_cache: List containing the cross-attention KV cache for each layer.

        Returns:
            logits: The final output logits (B, T, C * V), cast to float32.
        """
        _, _, num_channels_in = tgt_ids_BxTxC.shape
        assert num_channels_in == self.num_channels, "Input channels mismatch"

        # Embeddings
        x = None
        for i in range(self.num_channels):
            channel_tokens = tgt_ids_BxTxC[..., i]
            channel_embed = self.embeddings[i](channel_tokens)
            x = channel_embed if x is None else x + channel_embed

        # Apply dropout if not deterministic
        if not deterministic:
            x = self.dropout(x)

        # Process through each decoder layer
        for i, layer in enumerate(self.layers):
            x = layer(
                x,
                encoder_out,
                tgt_positions=tgt_positions,
                src_positions=src_positions,
                deterministic=deterministic,
                self_attn_mask=self_attn_mask,
                cross_attn_mask=cross_attn_mask,
                self_attn_cache=self_attention_cache[i],
                cross_attn_cache=cross_attention_cache[i],
                prefill=True,
            )

        # Final Norm
        x = self.norm(x)
        logits_BxTxCxV = self.logits_dense(x)

        # Convert to float32 if needed
        if logits_BxTxCxV.dtype != mx.float32:
            logits_BxTxCxV = logits_BxTxCxV.astype(mx.float32)

        return logits_BxTxCxV

        Forward pass for the Decoder stack, managing KV caches.

        Args:
            tgt_ids_BxTxC: Target token IDs (B, T, C).
            encoder_out: Output from the encoder (B, S, E).
            tgt_positions: Positions for target sequence (B, T).
            src_positions: Positions for source sequence (B, S).
            deterministic: Disable dropout if True.
            self_attn_mask: Mask for self-attention.
            cross_attn_mask: Mask for cross-attention.
            self_attention_cache: List containing the self-attention KV cache for each layer.
            cross_attention_cache: List containing the cross-attention KV cache for each layer.

        Returns:
            logits: The final output logits (B, T, C * V), cast to float32.
        
mlx_audio/tts/models/dia/layers.py