    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, c)
        else:
            q = k = v = x
            q, k, v = self.qkv_proj(q, n, k, n, v, n, b, c)
            q, q_shape, k, k_shape, v, v_shape = self.qkv_pool(q, k, v, thw_shape)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)

        n = q.shape[2]
        x = (attn @ v).transpose(1, 2).reshape(b, n, c)
        x = self.proj(x)
        if self.dropout_rate > 0.0:
            x = self.proj_drop(x)
        return x, q_shape
Args:
x(torch.Tensor):
Input tensor.
thw_shape(List):
The shape of the input tensor (before flattening).
towhee/models/layers/multi_scale_attention.py