src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/prime_attn
class JukeboxAttention:    def prime_attn(self, query, key, value, sample):
        encoder_len = self._encoder_len
        key = key[:, :encoder_len]
        value = value[:, :encoder_len]
        return self.dense_attn(query, key, value, sample)