revert to -->

Author: AntonOresten

README.md

@@ -10,7 +10,7 @@
 
 Einops.jl is a Julia implementation of the [einops](https://einops.rocks) Python package, providing an elegant and intuitive notation for tensor operations. We currently implement `rearrange`, offering a unified way to perform Julia's `reshape` and `permutedims`, as well as `pack` and `unpack`, with plans for implementing `reduce` and `repeat`.
 
-The Python implementation uses strings to specify the operation, but that would be tricky to compile in Julia, so a string macro `@einops_str` is exported for parity, e.g. `einops"a 1 b c -> (c b) a"`, which expands to the form `(:a, 1, :b, :c,) => ((:c, :b), :a)`, allowing for compile-time awareness of dimensionalities, ensuring type stability.
+The Python implementation uses strings to specify the operation, but that would be tricky to compile in Julia, so a string macro `@einops_str` is exported for parity, e.g. `einops"a 1 b c -> (c b) a"`, which expands to the form `(:a, 1, :b, :c,) --> ((:c, :b), :a)`, allowing for compile-time awareness of dimensionalities, ensuring type stability.
 
 ## Operations
 
@@ -22,15 +22,15 @@ The `rearrange` combines reshaping and permutation operations into a single, exp
 julia> images = randn(32, 30, 40, 3); # batch, height, width, channel
 
 # reorder axes to "b c h w" format:
-julia> rearrange(images, (:b, :h, :w, :c) => (:b, :c, :h, :w)) |> size
+julia> rearrange(images, (:b, :h, :w, :c) --> (:b, :c, :h, :w)) |> size
 (32, 3, 30, 40)
 
 # flatten each image into a vector
-julia> rearrange(images, (:b, :h, :w, :c) => (:b, (:h, :w, :c))) |> size
+julia> rearrange(images, (:b, :h, :w, :c) --> (:b, (:h, :w, :c))) |> size
 (32, 3600)
 
 # split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
-julia> rearrange(images, (:b, (:h1, :h), (:w1, :w), :c) => ((:b, :h1, :w1), :h, :w, :c), h1=2, w1=2) |> size
+julia> rearrange(images, (:b, (:h1, :h), (:w1, :w), :c) --> ((:b, :h1, :w1), :h, :w, :c), h1=2, w1=2) |> size
 (128, 15, 20, 3)
 ```
 
@@ -41,7 +41,7 @@ The `reduce` function will allow for applying reduction operations (like `sum`,
 ```julia
 # Example (conceptual):
 x = randn(100, 32, 64)
-y = reduce(maximum, x, (:t, :b, :c) => (:b, :c)) # max-reduction on the first axis
+y = reduce(maximum, x, (:t, :b, :c) --> (:b, :c)) # max-reduction on the first axis
 ```
 
 ### `repeat` (Planned)
@@ -51,7 +51,7 @@ The `repeat` function will provide a concise way to repeat elements along existi
 ```julia
 # Example (conceptual):
 image = randn(30, 40)
-rgb_image = repeat(image, (:h, :w) => (:repeat, :h, :w), repeat=3)
+rgb_image = repeat(image, (:h, :w) --> (:repeat, :h, :w), repeat=3)
 ```
 
 ## Roadmap

src/Einops.jl

@@ -7,6 +7,7 @@ export ..
 # TODO: use TransmuteDims.jl
 
 include("utils.jl")
+export -->
 
 include("einops_str.jl")
 export @einops_str
@@ -19,7 +20,7 @@ export pack, unpack
 
 # TODO: implement reduce, repeat
 Base.reduce(f, x::AbstractArray, pattern::Pattern; context...) = error("Not implemented")
-Base.repeat(x, pattern::Pattern; context...) = error("Not implemented")
+Base.repeat(x::AbstractArray, pattern::Pattern; context...) = error("Not implemented")
 
 # TODO: einsum
 

src/einops_str.jl

@@ -4,7 +4,7 @@ function parse_pattern(pattern::AbstractString)
     lhs, rhs = strip.(split(pattern, "->"; limit = 2))
     lhs_axes = tokenize_side(lhs)
     rhs_axes = tokenize_side(rhs)
-    return Tuple(lhs_axes) => Tuple(rhs_axes)
+    return Tuple(lhs_axes) --> Tuple(rhs_axes)
 end
 
 function tokenize_side(side::AbstractString)
@@ -86,10 +86,10 @@ For parity with Python implementation.
 
 ```jldoctest
 julia> einops"a 1 b c -> (c b) a"
-(:a, 1, :b, :c) => ((:c, :b), :a)
+(:a, 1, :b, :c) --> ((:c, :b), :a)
 
 julia> einops"embed token (head batch) -> (embed head) token batch"
-(:embed, :token, (:head, :batch)) => ((:embed, :head), :token, :batch)
+(:embed, :token, (:head, :batch)) --> ((:embed, :head), :token, :batch)
 
 julia> einops"i j * k" # for packing
 (:i, :j, *, :k)

src/rearrange.jl

@@ -46,22 +46,22 @@ end
 reshape_out(x, ::Tuple{Vararg{Symbol}}) = x
 
 """
-    rearrange(x::AbstractArray, left => right)
+    rearrange(x::AbstractArray, left --> right)
 
-Rearrange the axes of `x` according to the pattern specified by `left => right`.
+Rearrange the axes of `x` according to the pattern specified by `left --> right`.
 
 Can always be expressed as a `reshape` + `permutedims` + `reshape`.
 
 # Examples
 
 ```jldoctest
-julia> rearrange(rand(2,3,5), (:a, :b, :c) => (:c, :b, :a)) |> size
+julia> rearrange(rand(2,3,5), (:a, :b, :c) --> (:c, :b, :a)) |> size
 (5, 3, 2)
 
 julia> permutedims(rand(2,3,5), (3,2,1)) |> size
 (5, 3, 2)
 
-julia> rearrange(rand(2,3,35), (:a, :b, (:c, :d)) => (:a, :d, (:c, :b)), c=5) |> size
+julia> rearrange(rand(2,3,35), (:a, :b, (:c, :d)) --> (:a, :d, (:c, :b)), c=5) |> size
 (2, 7, 15)
 
 julia> reshape(permutedims(reshape(rand(2,3,35), 2,3,5,7), (1,4,3,2)), 2,7,5*3) |> size

src/utils.jl

@@ -1,3 +1,10 @@
+struct Pattern{L,R} end
+(-->)(L, R) = Pattern{L, R}()
+Base.show(io::IO, ::Pattern{L,R}) where {L,R} = print(io, "$L --> $R")
+Base.iterate(::Pattern{L}) where L = (L, Val(:R))
+Base.iterate(::Pattern{<:Any,R}, ::Val{:R}) where R = (R, nothing)
+Base.iterate(::Pattern, ::Nothing) = nothing
+
 function permutation_mapping(left::NTuple{N,T}, right::NTuple{N,T}) where {N,T}
     perm::Vector{Int} = findfirst.(isequal.([right...]), Ref([left...]))
     return ntuple(i -> perm[i], Val(N))

test/runtests.jl

@@ -18,14 +18,14 @@ using Test, Statistics
     @testset "einops string tokenization" begin
 
         @testset "arrow pattern" begin
-            @test einops"a b c -> a (c b)" == ((:a, :b, :c) => (:a, (:c, :b)))
-            @test einops"a b c -> a(c b)" == ((:a, :b, :c) => (:a, (:c, :b)))
-            @test einops"a b 1 -> a 1 b" == ((:a, :b, 1) => (:a, 1, :b))
-            @test einops"a b () -> a () b" == ((:a, :b, ()) => (:a, (), :b))
-            @test einops"a b()->a()b" == ((:a, :b, ()) => (:a, (), :b))
-            @test einops"b ... -> b ..." == ((:b, ..) => (:b, ..))
-            @test einops"->" == (() => ())
-            @test einops"-> 1" == (() => (1,))
+            @test einops"a b c -> a (c b)" == ((:a, :b, :c) --> (:a, (:c, :b)))
+            @test einops"a b c -> a(c b)" == ((:a, :b, :c) --> (:a, (:c, :b)))
+            @test einops"a b 1 -> a 1 b" == ((:a, :b, 1) --> (:a, 1, :b))
+            @test einops"a b () -> a () b" == ((:a, :b, ()) --> (:a, (), :b))
+            @test einops"a b()->a()b" == ((:a, :b, ()) --> (:a, (), :b))
+            @test einops"b ... -> b ..." == ((:b, ..) --> (:b, ..))
+            @test einops"->" == (() --> ())
+            @test einops"-> 1" == (() --> (1,))
             @test_throws "'.'" Einops.parse_pattern("-> .")
             @test_throws "'('" Einops.parse_pattern("-> (")
             @test_throws "')'" Einops.parse_pattern("-> )")
@@ -43,36 +43,36 @@ using Test, Statistics
     @testset "rearrange" begin
 
         x = rand(2,3,5)
-        @test rearrange(x, (:a, :b, :c) => (:c, :b, :a)) == permutedims(x, (3,2,1))
-        @test rearrange(x, (:a, :b, :c) => (:a, (:c, :b))) == reshape(permutedims(x, (1,3,2)), 2,5*3)
-        @test rearrange(x, (:first, :second, :third) => (:third, :second, :first)) == rearrange(x, (:a, :b, :c) => (:c, :b, :a))
-        @test_throws "Input length" rearrange(x, (:a, (:b, :c)) => (:c, :b, :a))
-        @test_throws ["Set of", "does not match"] rearrange(x, (:a, :b, :c) => (:a, :b, :a))
-        @test_throws ["Left names", "not unique"] rearrange(x, (:a, :a, :b) => (:a, :b))
-        @test_throws ["Right names", "not unique"] rearrange(x, (:a, :b, :c) => (:a, :b, :c, :a))
-        @test_throws "Invalid input dimension" rearrange(x, (:a, :b, 'c') => (:a, :b, :c))
-        @test_broken rearrange(x, (:a, :b, ..) => (:a, .., :b)) == rearrange(x, (:a, :b, :c) => (:a, :c, :b))
+        @test rearrange(x, (:a, :b, :c) --> (:c, :b, :a)) == permutedims(x, (3,2,1))
+        @test rearrange(x, (:a, :b, :c) --> (:a, (:c, :b))) == reshape(permutedims(x, (1,3,2)), 2,5*3)
+        @test rearrange(x, (:first, :second, :third) --> (:third, :second, :first)) == rearrange(x, (:a, :b, :c) --> (:c, :b, :a))
+        @test_throws "Input length" rearrange(x, (:a, (:b, :c)) --> (:c, :b, :a))
+        @test_throws ["Set of", "does not match"] rearrange(x, (:a, :b, :c) --> (:a, :b, :a))
+        @test_throws ["Left names", "not unique"] rearrange(x, (:a, :a, :b) --> (:a, :b))
+        @test_throws ["Right names", "not unique"] rearrange(x, (:a, :b, :c) --> (:a, :b, :c, :a))
+        @test_throws "Invalid input dimension" rearrange(x, (:a, :b, 'c') --> (:a, :b, :c))
+        @test_broken rearrange(x, (:a, :b, ..) --> (:a, .., :b)) == rearrange(x, (:a, :b, :c) --> (:a, :c, :b))
 
         x = reshape(rand(1)) # size (), length 1
-        @test rearrange(x, () => ()) == x
-        @test rearrange(x, () => (1,)) == reshape(x, 1)
+        @test rearrange(x, () --> ()) == x
+        @test rearrange(x, () --> (1,)) == reshape(x, 1)
 
         x = rand(2,3,5*7)
-        @test rearrange(x, (:a, :b, (:c, :d)) => (:a, :d, (:c, :b)), c=5) == reshape(permutedims(reshape(x, 2,3,5,7), (1,4,3,2)), 2,7,5*3)
+        @test rearrange(x, (:a, :b, (:c, :d)) --> (:a, :d, (:c, :b)), c=5) == reshape(permutedims(reshape(x, 2,3,5,7), (1,4,3,2)), 2,7,5*3)
 
         x = rand(2,3,5*7*11)
-        @test rearrange(x, (:a, :b, (:c, :d, :e)) => ((:a, :e), :d, (:c, :b)), c=5, d=7) == reshape(permutedims(reshape(x, 2,3,5,7,11), (1,5,4,3,2)), 2*11,7,5*3)
-        @test_throws "Unknown dimension sizes" rearrange(x, (:a, :b, (:c, :d, :e)) => (:a, :b, :c, :d, :e), c=5)
+        @test rearrange(x, (:a, :b, (:c, :d, :e)) --> ((:a, :e), :d, (:c, :b)), c=5, d=7) == reshape(permutedims(reshape(x, 2,3,5,7,11), (1,5,4,3,2)), 2*11,7,5*3)
+        @test_throws "Unknown dimension sizes" rearrange(x, (:a, :b, (:c, :d, :e)) --> (:a, :b, :c, :d, :e), c=5)
 
         x = rand(2,1,3)
-        @test rearrange(x, (:a, 1, :b) => (:a, :b)) == dropdims(x, dims=2)
-        @test_throws "Singleton dimension size is not 1" rearrange(x, (2, :a, :b) => (:a, :b))
-        @test_throws "Singleton dimension size is not 1" rearrange(x, (:a, :b, :c) => (:a, :b, :c, 2))
+        @test rearrange(x, (:a, 1, :b) --> (:a, :b)) == dropdims(x, dims=2)
+        @test_throws "Singleton dimension size is not 1" rearrange(x, (2, :a, :b) --> (:a, :b))
+        @test_throws "Singleton dimension size is not 1" rearrange(x, (:a, :b, :c) --> (:a, :b, :c, 2))
 
         x = rand(2,3)
-        @test rearrange(x, (:a, :b) => (:b, 1, :a)) == reshape(permutedims(x, (2,1)), 3,1,2)
-        @test rearrange(x, (:a, :b) => (:b, 1, 1, :a, 1)) == reshape(permutedims(x, (2,1)), 3,1,1,2,1)
-        @test rearrange(x, (:a, :b) => (:b, (), :a)) == rearrange(x, (:a, :b) => (:b, (), :a))
+        @test rearrange(x, (:a, :b) --> (:b, 1, :a)) == reshape(permutedims(x, (2,1)), 3,1,2)
+        @test rearrange(x, (:a, :b) --> (:b, 1, 1, :a, 1)) == reshape(permutedims(x, (2,1)), 3,1,1,2,1)
+        @test rearrange(x, (:a, :b) --> (:b, (), :a)) == rearrange(x, (:a, :b) --> (:b, (), :a))
 
         @testset "Python API reference parity" begin
             # see https://einops.rocks/api/rearrange/
@@ -106,6 +106,8 @@ using Test, Statistics
 
     @testset "reduce" begin
 
+        @test_throws "Not implemented" reduce(+, rand(2,3,4), einops"a b c -> b c")
+
         @testset "Python API reference parity" begin
             # see https://einops.rocks/api/reduce/
 
@@ -154,6 +156,8 @@ using Test, Statistics
 
     @testset "repeat" begin
 
+        @test_throws "Not implemented" repeat(rand(2,3,4), einops"a b c -> b c")
+
         @testset "Python API reference parity" begin
             # see https://einops.rocks/api/repeat/