Einops

It’s definitely better to just read doc

rearrange

• Change Sequence:

import einops
 
rearrange(img, "h w c -> w h c")
 
# Or: naming explicitly
 
rearrange(img, "width height color -> height width color")

• Composition of Dimensions

rearrange(images, "b h w c -> h (b w) c")

• Decomposition

# decomposition is the inverse process - represent an axis as a combination of new axes
# several decompositions possible, so b1=2 is to decompose 6 to b1=2 and b2=3
rearrange(images, "(b1 b2) h w c -> b1 b2 h w c ", b1=2)

• Order of axes matters

# The followings are very **different**
 
rearrange(ims, "b h w c -> h (b w) c")
rearrange(ims, "b h w c -> h (w b) c")
 

• Stack and concatenate

# rearrange can also take care of lists of arrays with the same shape
x = list(ims)
print(type(x), "with", len(x), "tensors of shape", x[0].shape)
 
# list of 6 tensor of shape (96,96,3)
 
# that's how we can stack inputs
# "list axis" becomes first ("b" in this case), and we left it there
rearrange(x, "b h w c -> b h w c").shape
 
# Shape now (6, 96, 96, 3)

reduce

• Reduce a dimension by simply making it disappear, using the operation you call

reduce(x, 'b h w c -> b h w', 'mean')

• Pooling operations:

# this is mean-pooling with 2x2 kernel
# image is split into 2x2 patches, each patch is averaged
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "mean", h2=2, w2=2)
 
# max-pooling is similar
# result is not as smooth as for mean-pooling
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "max", h2=2, w2=2)

repeat

• New dimension (axis)

# repeat along a new axis. New axis can be placed anywhere
repeat(ims[0], "h w c -> h new_axis w c", new_axis=5)

• Shortcut：

repeat（ims[0], "h w c -> h 5 w c"）

• Repeat along existing dimension

# repeat along w (existing axis)
repeat(ims[0], "h w c -> h (repeat w) c", repeat=3)
# Shortcut
repeat(ims[0]. "h w c -> h (3 w) c")

• Order also matters!