proposed optimized LAM (loop-less) #23
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR optimizes the LAM (Learned Adaptive Model) by removing explicit loops and implementing vectorized operations for better performance. The changes replace sequential frequency band processing with batch operations that process all bands simultaneously.
Key changes:
- Vectorized eigendecomposition and matrix operations across all frequency bands
- Removed explicit loops in encoding and decoding phases
- Maintained identical mathematical operations while improving computational efficiency
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| # Batch matrix multiply: [batch_size, freq_bands, Npx, N_ch] @ [batch_size, freq_bands, N_ch, N_ch] | ||
| # -> [batch_size, freq_bands, Npx, N_ch] | ||
| latent_x = torch.matmul(D_conj_T.unsqueeze(0), Vs) |
There was a problem hiding this comment.
Dimension mismatch in matrix multiplication. D_conj_T has shape [num_bands, Npx, Nch] and after unsqueeze becomes [1, num_bands, Npx, Nch], but Vs has shape [batch_size, freq_bands, N_ch, N_ch]. The matrix multiplication requires the inner dimensions to match, but here we have Nch vs N_ch which may not be aligned, and the frequency band dimensions need proper broadcasting.
| # Batch matrix multiply: [batch_size, freq_bands, Npx, N_ch] @ [batch_size, freq_bands, N_ch, N_ch] | |
| # -> [batch_size, freq_bands, Npx, N_ch] | |
| latent_x = torch.matmul(D_conj_T.unsqueeze(0), Vs) | |
| # Expand D_conj_T to match batch size for broadcasting | |
| # D_conj_T: [num_bands, Npx, Nch] -> [1, freq_bands, Npx, Nch] -> [batch_size, freq_bands, Npx, Nch] | |
| D_conj_T_expanded = D_conj_T.unsqueeze(0).expand(batch_size, -1, -1, -1) | |
| latent_x = torch.matmul(D_conj_T_expanded, Vs) |
No description provided.