src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/dense_attn
class JukeboxAttention:    def dense_attn(self, query, key, value, sample):
        query = self.split_heads(query)
        key = self.split_heads(key, is_key=True)
        value = self.split_heads(value)
        context_states = self._attn(query, key, value, sample)
        context_states = self.merge_heads(context_states)
        return context_states