    def forward(self, x, size):
        """forward function"""
        bb, n, cc = x.shape

        # Generate Q, K, V.
        qkv = (self.qkv(x).reshape(bb, n, 3, self.num_heads,
                                   cc // self.num_heads).permute(2, 0, 3, 1, 4))
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Factorized attention.
        k_softmax = k.softmax(dim=2)
        k_softmax_t_dot_v = einsum("b h n k, b h n v -> b h k v", k_softmax, v)
        factor_att = einsum("b h n k, b h k v -> b h n v", q,
                            k_softmax_t_dot_v)

        # Convolutional relative position encoding.
        crpe = self.crpe(q, v, size=size)

        # Merge and reshape.
        x = self.scale * factor_att + crpe
        x = x.transpose(1, 2).reshape(bb, n, cc)

        # Output projection.
        x = self.proj(x)
        x = self.proj_drop(x)

        return x
forward function
towhee/models/mpvit/mpvit.py