lucidrains / alphafold2

Hi, Luci:
Sorry, it's me again.
Here, I was confused by the definition of row-wise and col-wise attention.

Lines 208 to 209 in 586792d

    
           if self.row_attn: 
        
               w_x = rearrange(x, 'b h w d -> (b w) h d')

Lines 219 to 220 in 586792d

    
           if self.col_attn: 
        
               h_x = rearrange(x, 'b h w d -> (b h) w d')

Based on what I thought, the row w_x should be represented by (b h) w d, since once fetch one row, each row should have w(width) units.

So, maybe here need an inverse of the above definition?

The definition of row-wise attention and col-wise attention