src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/prime_qkv
class JukeboxAttention:    def prime_qkv(self, hidden_states, last_encoder_hidden_states=None, sample=False):
        curr_ctx = hidden_states.shape[1]
        if last_encoder_hidden_states is not None:
            raise TypeError("last_encoder_hidden_states should be None")
        query, key, value = hidden_states.chunk(3, dim=2)
        if sample:
            if self._cache_len() < self._encoder_len:
                self._append_cache(key, value)
            if self._cache_len() > self._encoder_len:
                self._slice_cache(0, self._encoder_len)
            key, value = self.cache["key"], self.cache["value"]
            self.sample_t += curr_ctx
        return query, key, value, sample