Transform input sequences into memory-organized chunks with capacity constraints. Processes input sequences by routing tokens to designated memory states according to routing_mask, sorts tokens by memory assignments, handles token truncation/padding based on memory capacity, and re
(
x: torch.Tensor,
routing_mask: torch.Tensor,
num_memories: int,
selected_memories: torch.Tensor,
attention_mask: torch.Tensor,
)
| 103 | |
| 104 | |
| 105 | def transform( |
| 106 | x: torch.Tensor, |
| 107 | routing_mask: torch.Tensor, |
| 108 | num_memories: int, |
| 109 | selected_memories: torch.Tensor, |
| 110 | attention_mask: torch.Tensor, |
| 111 | ): |
| 112 | ''' |
| 113 | Transform input sequences into memory-organized chunks with capacity constraints. |
| 114 | |
| 115 | Processes input sequences by routing tokens to designated memory states according to routing_mask, |
| 116 | sorts tokens by memory assignments, handles token truncation/padding based on memory capacity, |
| 117 | and returns memory-aligned tensors for parallel processing. |
| 118 | |
| 119 | Key operations: |
| 120 | 1. Expands input tensors when multiple memories are selected per token (top-k routing) |
| 121 | 2. Sorts tokens globally by (batch_idx, memory_idx) to group memory-assigned tokens |
| 122 | 3. Applies capacity-aware truncation (left-truncate oldest tokens when exceeding capacity) |
| 123 | 4. Pads memory chunks to uniform length for tensorization |
| 124 | |
| 125 | Args: |
| 126 | x: Input hidden states |
| 127 | Shape: (batch_size, seq_len, hidden_size) |
| 128 | routing_mask: Binary mask indicating active memory assignments |
| 129 | Shape: (batch_size, seq_len, num_memories) |
| 130 | num_memories: Total number of memories per batch |
| 131 | selected_memories: Memory indices assigned to each token. When using top-k routing, |
| 132 | this contains k memory indices per token (k >= 1) |
| 133 | Shape: (batch_size, seq_len) for k=1 or (batch_size, seq_len, topk) for k>1 |
| 134 | capacity: Scaling factor for memory capacity calculation. Actual capacity per memory is |
| 135 | ceil(seq_len * capacity), maintaining proportional capacity to sequence length |
| 136 | |
| 137 | Returns: |
| 138 | transformed_x: Memory-organized tensor with zero-padded capacity alignment |
| 139 | Shape: (num_memories, batch_size, capacity_len, hidden_size) |
| 140 | truncation_indices: Original indices used for gathering tokens after capacity truncation |
| 141 | Shape: (batch*num_memories, max_len) |
| 142 | sorted_indices: Sorting indices used to group tokens by memory assignments |
| 143 | Shape: (batch_size*seq_len*topk) |
| 144 | max_len: Maximum tokens per memory |
| 145 | mask: Boolean mask indicating valid (non-padded) positions in transformed_x |
| 146 | Shape: (batch*num_memories, max_len) |
| 147 | ''' |
| 148 | if selected_memories.dim() == 3: |
| 149 | # (batch, seq, topk) |
| 150 | topk = selected_memories.shape[2] |
| 151 | # x (batch, seq, hidden) |
| 152 | x = x.repeat_interleave(topk, dim=1) |
| 153 | # x (batch, seq * topk, hidden) |
| 154 | # (batch, seq, topk) |
| 155 | selected_memories = selected_memories.reshape(selected_memories.shape[0], -1) |
| 156 | # (batch, seq * topk) |
| 157 | |
| 158 | if attention_mask is not None: |
| 159 | attention_mask = attention_mask[:, -routing_mask.shape[1]:] |
| 160 | # mask out the masked tokens |
| 161 | routing_mask[attention_mask.bitwise_not().unsqueeze(-1).expand(-1, -1, num_memories)] = 0 |
| 162 |