src/transformers/models/jukebox/modeling_jukebox.py/JukeboxLayerStack/__init__
class JukeboxLayerStack:    def __init__(self, config, n_ctx):
        super().__init__()
        self.n_ctx = n_ctx
        self.width = config.hidden_size
        self.num_layers = config.num_layers
        self.blocks = config.blocks
        self.attention_pattern = config.attention_pattern
        if self.blocks is not None:
            self.block_ctx = n_ctx // self.blocks
        self.encoder_len = config.nb_relevant_lyric_tokens
        self.n_heads = config.n_heads

        # Orders of attn_func
        attention_pattern = ATTENTION_PATTERNS[self.attention_pattern]
        self._attn_mods = nn.ModuleList()
        for depth in range(self.num_layers):
            self._attn_mods.append(JukeboxBlock(config, n_ctx, attn_func=attention_pattern(depth)))

        self.saved_attn_weights = []