Are realrange and einsum really elegant

The conclusion is that , O.k. .

In terms of the amount of code , almost ：

# -*- coding:utf-8 -*-
import torch
from torch import nn
from torch import einsum
from einops import rearrange

class SimpleQKV(nn.Module):
    def __init__(self, dim, use_ein):
        super().__init__()
        self.proj = nn.Linear(dim, dim*3, bias=False)
        self.dim = dim
        self.scale = self.dim ** -0.5
        self.use_ein = use_ein
        torch.manual_seed(777) #  To make the weights the same , Easy to compare output 
        nn.init.xavier_uniform_(self.proj.weight)        

    def forward(self, x):
        n,c,h,w = x.shape
        #assert c==self.dim
        if (self.use_ein):
            x = rearrange(x, 'n c h w -> n (h w) c')
        else:
            x = x.permute(0,2,3,1).view(n, -1, c)
        qkv = self.proj(x)        
        q,k,v = qkv.chunk(chunks=3,dim=-1)        
        if (self.use_ein):
            attn = (einsum('n i c, n j c -> n i j', q, k) * self.scale).softmax(dim=-1)        
            v = einsum('n i j, n j c -> n i c', attn, v)
            output = rearrange(v, 'n (h w) c -> n c h w', h=h)
        else:            
            attn = (torch.matmul(q, k.transpose(1,2)) * self.scale).softmax(dim=-1)
            v = torch.matmul(attn, v)
            output = v.permute(0,2,1).view(n,c,h,w)
        
        return output

batch, chan, height, width = 1, 20, 7, 7
simple_qkv_ein = SimpleQKV(chan, True)
simple_qkv_noein = SimpleQKV(chan, False)

x = torch.randn(batch, chan, height, width, device='cpu')
out1 = simple_qkv_ein(x)
out2 = simple_qkv_noein(x)
assert(out1.equal(out2))

#  preservation onnx
simple_qkv_ein.eval()
onnx_filename = './simple_qkv_ein.onnx'
torch.onnx.export(simple_qkv_ein, x, onnx_filename,
                  input_names=['input'], output_names=['ouput'],
                  export_params=True, verbose=False, opset_version=12)

simple_qkv_noein.eval()
onnx_filename = './simple_qkv_noein.onnx'
torch.onnx.export(simple_qkv_noein, x, onnx_filename,
                  input_names=['input'], output_names=['ouput'],
                  export_params=True, verbose=False, opset_version=12)

print('save onnx succ.')

From saved onnx see （ after onnxsim Optimize ）, Also almost ：