src/transformers/models/jukebox/configuration_jukebox.py/enc_dec_with_lyrics
def enc_dec_with_lyrics(layer):
    if layer % 16 == 15:
        return _PrimePrimeDenseAttention[layer % 3]
    return _RawColumnPreviousRowAttention[layer % 3]