src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/forward
class JukeboxAttention:    def forward(self, hidden_states, last_encoder_hidden_states=None, sample=False):
        curr_ctx = hidden_states.shape[1]
        hidden_states = self.c_attn(hidden_states)
        query, key, value, sample = self.qkv(
            hidden_states, last_encoder_hidden_states=last_encoder_hidden_states, sample=sample
        )
        attention_scores = self.attn(query, key, value, sample)
        if attention_scores.shape[1] != curr_ctx:
            offset = self._offset(curr_ctx)
            attention_scores = attention_scores[:, offset : offset + curr_ctx, :].contiguous()
        attention_scores = self.c_proj(attention_scores)
        return self.resid_dropout(attention_scores)