Macro-based interface

[AntonOresten]

Aside from the @einops_str macro for Python syntax, an @einops macro could be used to reduce colons and commas in the input expression—although without commas, (a b) won't work, so [a b] could be used instead.

@einops a b c d -> [a b] c d

This is parsed as:

julia> :(@einops a b c d -> [a b] c d)
:(#= REPL[5]:1 =# @einops a b c (d->begin
              #= REPL[5]:1 =#
              [a b]
          end) c d)

so anything adjacent to -> would require special parsing.

Alternatively (or additionally), the @einops macro could act on calls, e.g. @einops rearrange(x, "a b c d -> (a b) c d") but this isn't really better than rearrange(x, einops"a b c d -> (a b) c d").

Another option is @rearrange(x, "a b c d -> (a b) c d") for near perfect parity with Python implementation (just an extra @).

I don't want this becoming full of macros though, with 5 different ways of calling the same thing.

I would potentially consider just shortening einops"..." to ein"..." like they do in OMEinsum.jl, but there's the risk of clashing.

[AntonOresten]

Another option is @rearrange(x, "a b c d -> (a b) c d") for near perfect parity with Python implementation (just an extra @).

The entire package should arguably use macros to move more work to macro-expansion time instead of runtime, despite the latter still somehow being type-stable and Reactant-compatible. With macros, it would just be a matter of translating the expression into the underlying reshapes, permutations, etc. I don't have a clear idea of how this would look, but it seems like it would be the Julian way.

[AntonOresten]

A macro-based implementation would likely also solve #4.

[AntonOresten]

Macros would also allow for interpolation syntax, such that you could write:

@rearrange(x, (($a, b), c) --> (a, b, c))

instead of:

rearrange(x, ((:a, :b), :c) --> (:a, :b, :c), a=a)

This is also good because macros don't natively accept keyword arguments anyway (as it counts as expressions), but requiring the name to match to a variable might be silly.

It's tricky because as we see in above comments, the API can sensibly be defined in a myriad of ways. The current is the one that might make most sense intuitively, and follows closely the original einops python package syntax, but simultaneously Julia's macros would solve a lot of problems, including making code shorter and moving all work to macro expansion time (big plus).

I don't want to arbitrarily change the API. The package works, and provides some simple functionality/syntax that the ecosystem lacked in some ways. I should only consider making changes if users of the package are interested.

[AntonOresten]

My fear is that ellipses will be tricky in a macro version, since the type of the array (and thus its dimensionality) isn't known at macro-expansion time.
Might be solvable by slicing the array size (size(x)[...]) with some range that depends on the ellipsis placement.

[AntonOresten]

There is an argument to be made that macros can be confusing.
I sometimes feel viscerally repelled when I see a macro, especially when they have lots of rules and custom syntax. The current rearrange function is elegantly intuitive, albeit slightly more verbose compared to the Python syntax.

[AntonOresten]

Maybe e.g. rearrange can be a generated function, that changes the returned expression based on the pattern. This still requires a very different approach, but unlike a macro-based interface, it could possibly be made non-breaking.

[AntonOresten]

An interesting calling pattern I came up with, that might not be possibly with other approaches. It squeezes a repeated operation into a single line:

q, k = rearrange.((q, k), ((:d, :h) ..) --> (:d, :h, ..))

Should still work with a generated function approach.

Note that support for this was added in v0.1.7

[AntonOresten]

[AntonOresten]

e.g. @rearrange would basically be a non-specializing inline version of the current rearrange. Not sure about the trade-offs. Generated functions are probably more expensive? But at the same time, macros prevent you from doing broadcasting like:

q, k, v = rearrange.((q, k, v), einops"(d h) l ... -> d l h ..."; d=layer.head_dim)

which is an insanely cool trick that allowed me to do a very terse grouped-query attention implementation. So this is one of the arguments against macros.

[AntonOresten]

I guess there's nothing preventing us from having both... although if it's inlining we're after, then see: https://groups.google.com/g/julia-users/c/iZhWp5YSqmc?pli=1. I guess the main benefit of an explicit macro interface is that they can be expanded to understand what's going on. current workaround is to dev the package, remove @generated, and write e.g. context = typeof(context) since the args are all the types of the actual arguments, outside the returned expression within generated functions.