Problem: In one of my programs, I need to calculate a matrix multiplication A @ B where both are of size N by N for considerably large N. I'm conjecturing that approximating this product by using band_matrix(A, width) @ B could just suffice the needs, where band_matrix(A, width) denotes a band matrix part of A with width width. For example, width = 0 gives the diagonal matrix with diagonal elements taken from A and width = 1 gives the tridiagonal matrix taken in a similar manner. My try: I'm trying to extract the tridiagonal matrix, for instance, in the following way: # Step 1: Extract the main diagonal main_diag = torch.diagonal(A, dim1=-2, dim2=-1) # Shape: [d1, d2, N] # Step 2: Extract the upper diagonal (offset=1) upper_diag = torch.diagonal(A, offset=1, dim1=-2, dim2=-1) # Shape: [d1, d2, N-1] # Step 3: Extract the lower diagonal (offset=-1) lower_diag = torch.diagonal(A, offset=-1, dim1=-2, dim2=-1) # Shape: [d1, d2, N-1] # Step 4: Reconstruct the tridiagonal matrix # Main diagonal tridiag = torch.diag_embed(main_diag) # Shape: [d1, d2, N, N] # Upper diagonal (shift the values to create the first upper diagonal) tridiag += torch.diag_embed(upper_diag, offset=1) # Lower diagonal (shift the values to create the first lower diagonal) tridiag += torch.diag_embed(lower_diag, offset=-1) but I'm not sure if tridiag @ B would be much more efficient than the original A @ B or just the same complexity since Torch may not know the specific structure to tridiag. In theory, computation with a tridiagonal matrix should be N times faster. Any help with understanding PyTorch's behaviour in this type of scenario or implementing some alternative GPU optimized approaches would be greatly appreciated. Continue reading...
