Rearrange to reshape and permute?

[camenduru]

Hi I am learning how to use einops with https://einops.rocks/pytorch-examples.html

I have a question how to convert this q, k, v = map(lambda t: rearrange(t, 'b c h w -> b (h w) c'), (q, k, v)) rearrange() function to reshape() and permute()

[arogozhnikov]

Hi @camenduru

You can think of it in steps:

# first transpose: b c h w -> b h w c
x = x.transpose([0, 2, 3, 1])

# then compose two axes into one: b h w c-> b (h w) c
b, h, w, c = x.shape
x = x.reshape([b,  h * w, c])

stackoverflow is the right place for such questions

[camenduru]

thank you and sorry next time I will use stackoverflow.