    def forward(self,
                query,
                key,
                value,
                key_padding_mask=None,
                incremental_state=None,
                attn_mask=None):
        """
        Inputs of forward function
            query: [target length, batch size, embed dim]
            key: [sequence length, batch size, embed dim]
            value: [sequence length, batch size, embed dim]
            key_padding_mask: if True, mask padding based on batch size
            incremental_state: if provided, previous time steps are cashed
            need_weights: output attn_output_weights
            static_kv: key and value are static

        Outputs of forward function
            attn_output: [target length, batch size, embed dim]
            attn_output_weights: [batch size, target length, sequence length]
        """
        q_shape = query.shape
        src_shape = key.shape
        q = self._in_proj_q(query)
        k = self._in_proj_k(key)
        v = self._in_proj_v(value)
        q *= self.scaling
        # q = paddle.transpose(
        #     paddle.reshape(
        #         q, [q_shape[0], q_shape[1], self.num_heads, self.head_dim]),
        #     [1, 2, 0, 3])
        q = torch.reshape(q, (q_shape[0], q_shape[1], self.num_heads, self.head_dim))
        q = q.permute(1, 2, 0, 3)
        # k = paddle.transpose(
        #     paddle.reshape(
        #         k, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]),
        #     [1, 2, 0, 3])
        k = torch.reshape(k, (src_shape[0], q_shape[1], self.num_heads, self.head_dim))
        k = k.permute(1, 2, 0, 3)
        # v = paddle.transpose(
        #     paddle.reshape(
        #         v, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]),
        #     [1, 2, 0, 3])
        v = torch.reshape(v, (src_shape[0], q_shape[1], self.num_heads, self.head_dim))
        v = v.permute(1, 2, 0, 3)
        if key_padding_mask is not None:
            assert key_padding_mask.shape[0] == q_shape[1]
            assert key_padding_mask.shape[1] == src_shape[0]
        attn_output_weights = torch.matmul(q,
                                            k.permute(0, 1, 3, 2))
        if attn_mask is not None:
            attn_mask = torch.unsqueeze(torch.unsqueeze(attn_mask, 0), 0)
            attn_output_weights += attn_mask
        if key_padding_mask is not None:
            attn_output_weights = torch.reshape(
                attn_output_weights,
                [q_shape[1], self.num_heads, q_shape[0], src_shape[0]])
            key = torch.unsqueeze(torch.unsqueeze(key_padding_mask, 1), 2)
            key = key.type(torch.float32)
            y = torch.full(
                size=key.shape, fill_value=float("-Inf"), dtype=torch.float32)
            y = torch.where(key == 0., key, y)
            attn_output_weights += y
        attn_output_weights = F.softmax(
            attn_output_weights.type(torch.float32),
            dim=-1,
            dtype=torch.float32 if attn_output_weights.dtype == torch.float16
            else attn_output_weights.dtype)
        attn_output_weights = F.dropout(
            attn_output_weights, p=self.dropout, training=self.training)

        attn_output = torch.matmul(attn_output_weights, v)
        attn_output = torch.reshape(
        attn_output.permute(2, 0, 1, 3),
            [q_shape[0], q_shape[1], self.embed_dim])
        attn_output = self.out_proj(attn_output)

        return attn_output
Inputs of forward function
query: [target length, batch size, embed dim]
key: [sequence length, batch size, embed dim]
value: [sequence length, batch size, embed dim]
key_padding_mask: if True, mask padding based on batch size
incremental_state: if provided, previous time steps are cashed
need_weights: output attn_output_weights
static_kv: key and value are static

Outputs of forward function
attn_output: [target length, batch size, embed dim]
attn_output_weights: [batch size, target length, sequence length]
pytorchocr/modeling/heads/multiheadAttention.py