def _create_attn_mask(
        self,
        q_padding_mask_1d: mx.array,
        k_padding_mask_1d: mx.array,
        is_causal: bool = False,
    ) -> mx.array:
        """
        Creates the attention mask (self or cross) mimicking JAX segment ID logic.
        """
        B1, Tq = q_padding_mask_1d.shape
        B2, Tk = k_padding_mask_1d.shape
        assert B1 == B2, "Query and key batch dimensions must match"

        p_mask_q = mx.expand_dims(q_padding_mask_1d, 2)  # Shape [B, Tq, 1]
        p_mask_k = mx.expand_dims(k_padding_mask_1d, 1)  # Shape [B, 1, Tk]

        # Condition A: Non-padding query attends to non-padding key
        non_pad_attends_non_pad = mx.logical_and(
            p_mask_q, p_mask_k
        )  # Shape [B, Tq, Tk]

        # Condition B: Padding query attends to padding key
        pad_attends_pad = mx.logical_and(
            mx.logical_not(p_mask_q), mx.logical_not(p_mask_k)
        )  # Shape [B, Tq, Tk]

        # Combine: True if padding status is compatible (both non-pad OR both pad)
        # This implementation follows Jax TPU splash attention kernel
        mask = mx.logical_or(
            non_pad_attends_non_pad, pad_attends_pad
        )  # Shape [B, Tq, Tk]

        if is_causal:
            # Ensure causality for self-attention (Tq == Tk)
            assert (
                Tq == Tk
            ), "Causal mask requires query and key sequence lengths to be equal"
            # Standard lower-triangular causal mask (True means allow)
            causal_mask_2d = mx.tril(
                mx.ones((Tq, Tk), dtype=mx.bool_)
            )  # Shape [Tq, Tk]
            causal_mask = mx.logical_and(mask, causal_mask_2d)  # Shape [B, Tq, Tk]
            return mx.expand_dims(
                causal_mask, 1
            )  # Shape [B, 1, Tq, Tk] for broadcasting across heads
        else:
            # For cross-attention or non-causal self-attention
            return mx.expand_dims(
                mask, 1
            )  # Shape [B, 1, Tq, Tk] for broadcasting across heads

        Creates the attention mask (self or cross) mimicking JAX segment ID logic.
        
mlx_audio/tts/models/dia/dia.py