src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/_offset
class JukeboxAttention:    def _offset(self, curr_ctx):
        if self.attn_func == "dense_attn":
            return 0
        return (self.sample_t - curr_ctx) % self.block_ctx