Alias ellipsis to a star for rearrange

[knyazer]

Hey, guys, the library is great!

But, there is a small non-intuitive thing that I consistently stumble upon (I mean, twice is consistent, right?). It is that for rearrange one have to use the ellipsis instead of a star to select the rest of dimensions. It seems a bit nicer to have an alias of ellipsis to a star, e.g.

rearrange('batch_dim_1 batch_dim_2 ... -> batch_dim_2 batch_dim_1 ...')
# proposal: rearrange('batch_dim_1 batch_dim_2 * -> batch_dim_2 batch_dim_1 *')

Probably I find the star notation more sensible, since I use jaxtyping a lot, and it uses the star notation for this. If you have a good argument for why the star should not be equivalent to the ellipsis, I am happy to just alias this in jaxtyping, so that the dimension selection convention is the same between the two: I guess many people use both of the libraries at the same time (at least I do :) )

P.S. I am new-ish with einops, so I might not know that this problem was already discussed a million times. In this case, just point me to the relevant discussion :smile:

[arogozhnikov]

Hi @knyazer

this problem was already discussed a million times

you'll be surprised, but it is the first time it is discussed in repo :)

First, how did we end up in where we are (that's how I reconstruct history) in chronological order

• f(*x) for passing multiple args
• ellipsis (...) used in numpy indexing x[..., 0]
• numpy.einsum decided to use ... for a group of unnamed axes
• python3 introduces star into unpacking multiple arguments [a, *b] = f(x)
• einops follows einsum with using ellipsis in rearrange
• einops uses * in unpack to designate special position for one-or-many axes
• jaxtyping uses * to designate named or unnamed group of axes

so multiple axes historically was always ellipsis, but introduction of star-unpacking makes it tempting to use star and align with less array-specific part of python.

Two considerations:

• star has a bit different meaning in unpack, not a blocker, but e.g. #253 would be blocked
• mixture of styles is a bit ugly, and there is a lot of einsum/einops code using ellipsis

cc: @patrick-kidger

[knyazer]

Huh, interesting.

---

Next part is mostly outdated due to Patrick's comment, I had some misconceptions about jaxtyping…

Feel free to ignore it :)

---

I read about what the star does in einops, and it is equivalent to a underscore in jaxtyping! And ellipsis in einops is equivalent to a star in jaxtyping.

The only way I see to synchronize the APIs without breaking backwards compatibility, is to:

1. Alias ellipsis to a star in jaxtyping, don't deprecate star, make ellipsis a "recommended" way.
2. Alias underscore to a star in einops, prefer underscore over star.

Ta-dah! Now we can use underscore to mark a single arbitrary dimension, and an ellipsis to mark multiple arbitrary axes in both libraries.

But, this is a very significant API change. So, I don't know how you feel about that. However, I think it is useful: the difference between star and ellipsis is indeed confusing (to me). And underscore is probably a more pythonic way to show a single arbitrary axis (if you think about axis sizes as variables).

Whether star or ellipsis more pythonic for multiple dims is up to a debate, but only ellipsis does not break backwards compatibility, so there is not really a choice.

WDYT?

[patrick-kidger]

Hey folks!

So I think we have one baked-in inconsistency already. In einops, ... is (IIUC) a "named" group of dimensions: it must be the same on both sides. In jaxtyping, it is an anything-goes, and in particular can represent different groups of axes across its different usages. So if consistency was a goal then I think we'd both have to recommend avoiding ... entirely.

We'd then find some way of both libraries using * in compatible ways? FWIW a plain * on its own is not currently valid notation in jaxtyping, so we have an opportunity there. I'm not sure how important consistency really is between each other, but I'd be happy to tweak things if it's possible!

Frankly this inconsistency is my mistake. when I came up with the notation for jaxtyping, I completely forgot to check what notation einops was already using... 😅