einops

This is a very OP library. Really learned at Dyna Robotics.

Really import to see that these two are not the same:

(Pdb) import einops
(Pdb) x = torch.tensor([[1, 2], [3, 4]])
(Pdb) x
tensor([[1, 2],
        [3, 4]])
(Pdb) einops.repeat(x, 'b d -> (b repeat) d', repeat=3)
tensor([[1, 2],
        [1, 2],
        [1, 2],
        [3, 4],
        [3, 4],
        [3, 4]])
(Pdb) einops.repeat(x, 'b d -> (repeat b) d', repeat=3)
tensor([[1, 2],
        [3, 4],
        [1, 2],
        [3, 4],
        [1, 2],
        [3, 4]])

For arbitrary shape, you can do ... :) Very convenient

Other questions

if you do:

torch.cat([x] * repeat, dim=0)

That is functionally identical to:

einops.repeat(x, 'b d -> (repeat b) d', repeat=repeat)

• Notice that repeat needs to go in front of b

x.repeat(repeat, 1)