src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/summary_spread_attn
class JukeboxAttention:    def summary_spread_attn(self, query, key, value, sample):
        blocks = self.blocks
        spread = self.spread

        batch_size, seq_len, embed_dim = value.shape  # For sample, query_len= 1, key_len = value_len = sample_t
        if sample:
            raise NotImplementedError
        else:
            key = key.view(batch_size, blocks, seq_len // blocks, embed_dim)[:, :-1, -spread:, :]
            key = torch.nn.functional.pad(key, (0, 0, 0, 0, 1, 0)).contiguous()
            key = key.view(batch_size, blocks * spread, embed_dim)

            value = value.view(batch_size, blocks, seq_len // blocks, embed_dim)[:, :-1, -spread:, :]
            value = torch.nn.functional.pad(value, (0, 0, 0, 0, 1, 0)).contiguous()
            value = value.view(batch_size, blocks * spread, embed_dim)

            return self.dense_attn(query, key, value, sample).view(batch_size, seq_len, embed_dim)