def __call__(self, q, k, v, mask=None):
        """
        einstein notation
        b - batch
        h - heads
        n, i, j - sequence length (base sequence length, source, target)
        d - feature dimension
        """
        n = q.shape[-2]

        scale = q.shape[-1] ** -0.5

        # Handle different dimensions for k and v
        kv_einsum_eq = "b j d" if k.ndim == 3 else "b h j d"

        # similarity
        if k.ndim == 3:
            k = mx.expand_dims(k, axis=1)
            k = mx.broadcast_to(k, q.shape)

        if v.ndim == 3:
            v = mx.expand_dims(v, axis=1)
            v = mx.broadcast_to(v, q.shape[:-1] + (v.shape[-1],))

        # q: [b h i d], k: [b h j d]
        sim = mx.matmul(q, mx.transpose(k, (0, 1, 3, 2))) * scale

        # key padding mask
        if exists(mask):
            mask = mx.reshape(mask, (mask.shape[0], 1, 1, mask.shape[1]))
            sim = mx.where(mask, sim, -1e9)

        # causal mask
        if self.causal:
            causal_mask = self.get_mask(n)
            sim = mx.where(causal_mask, -1e9, sim)

        # attention
        attn = mx.softmax(sim, axis=-1)

        if self.dropout > 0 and self.training:
            attn = self.attn_dropout(attn)

        # aggregate values
        out = mx.matmul(attn, v)

        return out

        einstein notation
        b - batch
        h - heads
        n, i, j - sequence length (base sequence length, source, target)
        d - feature dimension
        
mlx_audio/tts/models/spark/modules/speaker/perceiver_encoder.py