[Feature suggestion] Operations on elements of a dimension

Hello, I’m just throwing an idea, I’m not sure it fits in the scope of Einops, and it will probably require a lot of work, but I think it would be useful:

What about allowing to manipulate elements along a dimension, like for example the r,g,b channels of an image ? I could imagine a syntax that would look like this, with everything inside brackets referring to elements rather than dimensions:

# reorder color channels rgb to bgr OpenCV-style:
rearrange(imgs, 'batch [r g b] h w -> batch [b g r] h w')

Extending on the existing syntax, grouping elements would look like this:

# reorder color channels rgb to brg:
rearrange(imgs, 'batch [rg b] h w -> batch [b rg] h w', rg=2)

This could allow to drop elements:

# remove alpha channel:
rearrange(imgs, 'batch [rgb a] h w -> batch [rgb] h w', rgb=3)

[remisphere]

For the h[top bottom] vs [top bottom] part, I was trying to anticipate cases where you would also want to rearrange dimensions, like

rearrange(imgs, 'c[rgb a] h[top bottom] -> h[top] c[rgb]', rgb=3, top=128)

because I am not sure how hard it is to recognise dimensions based on content, and it may arguably be easier to see what the dimension represents.

Indeed, it may make sense to put that functionality in a new function to keep readability. But then how do we show in an obvious way on which dimension to operate ? Maybe something like this (for working on multiple dimensions at a time):

# maybe `reorder` is not the best name, I think it is too close to `rearrange` in its meaning.
reorder(imgs, 'batch [rg b a -> b rg] [top bottom -> top] w', rg=2, b=1, top=128, batch=10)

while forbidding to rearrange dimensions at the same time to keep the code clear and concerns separated ? with an -> operator for potentially each dimension. (or we can keep the previous notation with a single arrow)

If the function is made to only work on one axis at a time, there should be a clear way of stating which one. I don’t see a way of doing it without writing at least once all dimensions, with something like this:

reorder(imgs, 'batch [c] h w: rgb a -> rgb', rgb=3)

[arogozhnikov]

einops.pack and einops.unpack are the current way to manipulate a single dimension.

They immediately cover chop and concatenation from this suggestion. As for reordering, it can be implemented with a combination unpack+pack.

You can find some closer-to-application examples here