class_docstr: Implementation of a multiscale attention block. Compare to a conventional attention
  block, a multiscale attention block optionally supports pooling (either
  before or after qkv projection). If pooling is not used, a multiscale attention
  block is equivalent to a conventional attention block.
  
  ::
                                 Input
                                   |
                  |----------------|-----------------|
                  ↓                ↓                 ↓
                Linear           Linear            Linear
                  &                &                 &
               Pool (Q)         Pool (K)          Pool (V)
                  → -------------- ←                 |
                           ↓                         |
                     MatMul & Scale                  |
                           ↓                         |
                        Softmax                      |
                           → ----------------------- ←
                                       ↓
                                 MatMul & Scale
                                       ↓
                                    DropOut
docstr: Args:
      x (torch.Tensor): Input tensor.
      thw_shape (List): The shape of the input tensor (before flattening).
docstr_tok: ["Args",":","x","(","torch.Tensor",")",":","Input","tensor",".","thw_shape","(","List",")",":","The","shape","of","the","input","tensor","(","before","flattening",")","."]
code: def forward(
      self, x: torch.Tensor, thw_shape: List[int]
  ) -> Tuple[torch.Tensor, List[int]]:
      """
      Args:
          x (torch.Tensor): Input tensor.
          thw_shape (List): The shape of the input tensor (before flattening).
      """
  
      B, N, C = x.shape
      qkv = (
          self.qkv(x)
          .reshape(B, N, 3, self.num_heads, -1)
          .permute(2, 0, 3, 1, 4)
      )
  
      q, k, v = qkv[0], qkv[1], qkv[2]
      q, q_shape, k, k_shape, v, v_shape = self._qkv_pool(
          q, k, v, thw_shape
      )
  
      attn = (q * self.scale) @ k.transpose(-2, -1)
      attn = attn.softmax(dim=-1)
  
      N = q.shape[2]
  
      if self.residual_pool:
          x = (attn @ v + q).transpose(1, 2).reshape(B, N, C)
      else:
          x = (attn @ v).transpose(1, 2).reshape(B, N, C)
  
      x = self.proj(x)
      return x, q_shape
  
imports: from typing import List, Optional, Tuple
  import numpy
  import torch
  import torch.fx
  import torch.nn
  from torch.nn.common_types import _size_3_t
  from .drop_path import DropPath
label_desc: Use image node_id 6 for calling the MultiScaleAttention obj's underlying member method code with example usage: obj.forward(x, thw_shape) and returns: x,  q_shape