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
             [dim expand]     [dim expand]      [dim expand]
                  &                &                 &
               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
      if self.pool_first:
          x = x.reshape(
              B, N, self.num_heads, C // self.num_heads
          ).permute(0, 2, 1, 3)
          q = k = v = x
          q, q_shape, k, k_shape, v, v_shape = self._qkv_pool(
              q, k, v, thw_shape
          )
          q_N, k_N, v_N = self._get_qkv_length(
              q_shape, k_shape, v_shape
          )
          q, k, v = self._reshape_qkv_to_seq(
              q, k, v, q_N, v_N, k_N, B, C
          )
          q, k, v = self._qkv_proj(
              q, q_N, k, k_N, v, v_N, B, self.dim_out
          )
      else:
          if self.separate_qkv:
              q = k = v = x
              q, k, v = self._qkv_proj(
                  q, N, k, N, v, N, B, self.dim_out
              )
          else:
              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, -1, self.dim_out)
          )
      else:
          x = (attn @ v).transpose(1, 2).reshape(B, -1, self.dim_out)
  
      x = self.proj(x)
      if self.dropout_rate > 0.0:
          x = self.proj_drop(x)
      return x, q_shape
  
imports: from typing import Callable, List, Optional, Tuple
  import numpy
  import torch
  import torch.nn
  from torch.nn.common_types import _size_3_t
  from .drop_path import DropPath
label_desc: Use image node_id 7 for calling the MultiScaleAttention obj's underlying member method code with example usage: obj.forward(x, thw_shape) and returns: x,  q_shape