stow away --> notation in favor of =>, docstrings

Author: AntonOresten

README.md

@@ -8,9 +8,9 @@
 > [!WARNING]  
 > This package is still under development, and does not implement all of the features of the original (see [Roadmap](#Roadmap)).
 
-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`, with plans for also implementing `reduce` and `repeat`.
+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)` where `-->` is an operator that creates an `Einops.Pattern{(:a, 1, :b, :c), ((:c, :b), :a)}`, allowing for compile-time awareness of dimensionalities and permutations—this is not yet taken advantage of, since the tuple types are sufficient for at least ensuring type stability (see [Roadmap](#Roadmap)).
+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
@@ -63,7 +63,7 @@ rgb_image = repeat(image, (:h, :w) --> (:repeat, :h, :w), repeat=3)
 *   [ ] Implement `repeat`.
 *   [ ] Explore integration with `PermutedDimsArray` or `TransmuteDims.jl` for lazy and statically inferrable permutations.
 *   [ ] Implement `einsum` (or wrap existing implementation) 
-*   [ ] Implement `pack`, and `unpack`.
+*   [x] Implement `pack`, and `unpack`.
 
 ## Contributing
 

src/Einops.jl

@@ -6,15 +6,7 @@ export ..
 
 # TODO: use TransmuteDims.jl
 
-struct Pattern{L,R} end
-(-->)(left, right) = Pattern{left, right}()
-Base.show(io::IO, ::Pattern{L,R}) where {L,R} = print(io, "$L --> $R")
-Base.iterate(::Pattern{L}) where L = (L, Val(:right))
-Base.iterate(::Pattern{<:Any,R}, ::Val{:right}) where R = (R, nothing)
-Base.iterate(::Pattern, ::Nothing) = nothing
-
 include("utils.jl")
-export -->
 
 include("einops_str.jl")
 export @einops_str

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/pack_unpack.jl

@@ -9,23 +9,68 @@ size_after_wildcard(dims::Dims, pattern::PackPattern) = dims[end-length(pattern)
 size_wildcard(dims::Dims, pattern::PackPattern) = dims[begin+length(size_before_wildcard(dims, pattern)):end-length(size_after_wildcard(dims, pattern))]
 packed_size(dims::Dims, pattern::PackPattern) = (size_before_wildcard(dims, pattern)..., prod(size_wildcard(dims, pattern)), size_after_wildcard(dims, pattern)...)
 
-function pack(unpacked_arrays::Vector, pattern::PackPattern{N}) where N
+"""
+    pack(unpacked_arrays, pattern)
+
+Pack a vector of arrays into a single array according to the pattern.
+
+# Examples
+
+```jldoctest
+julia> inputs = [rand(2,3,5), rand(2,3,7,5), rand(2,3,7,9,5)]
+
+julia> packed_array, packed_shapes = pack(inputs, (:i, :j, *, :k));
+
+julia> size(packed_array)
+(2, 3, 71, 5)
+
+julia> packed_shapes
+3-element Vector{NTuple{N, Int64} where N}:
+ ()
+ (7,)
+ (7, 9)
+"""
+function pack(unpacked_arrays, pattern::PackPattern{N}) where N
     checkpacking(pattern)
     reshaped_arrays = [reshape(A, packed_size(size(A), pattern)::Dims{N})::AbstractArray{<:Any,N} for A in unpacked_arrays]
     concatenated_array::AbstractArray{<:Any,N} = cat(reshaped_arrays..., dims=find_wildcard(pattern))
-    packed_shapes = Dims[size_wildcard(size(A), pattern) for A in unpacked_arrays]
+    packed_shapes = Dims[size_wildcard(size(unpacked_array), pattern) for unpacked_array in unpacked_arrays]
     return concatenated_array, packed_shapes
 end
 
 splice(a::Dims, i::Int, r::Dims) = (a[1:i-1]..., r..., a[i+1:end]...)
 
-# FIXME: results in a reshape of a view ... so should we collect?
+# FIXME: result is a reshape of a view ... so we should collect?
+"""
+    unpack(packed_array, packed_shapes, pattern)
+
+Unpack a single array into a vector of arrays according to the pattern.
+
+# Examples
+
+```jldoctest
+julia> inputs = [rand(2,3,5), rand(2,3,7,5), rand(2,3,7,9,5)];
+
+julia> inputs == unpack(pack(inputs, (:i, :j, *, :k))..., (:i, :j, *, :k))
+true
+
+julia> packed_array = rand(2,3,16);
+
+julia> packed_shapes = [(), (7,), (4, 2)];
+
+julia> unpack(packed_array, packed_shapes, (:i, :j, *)) .|> size
+3-element Vector{Tuple{Int64, Int64, Vararg{Int64}}}:
+ (2, 3)
+ (2, 3, 7)
+ (2, 3, 4, 2)
+```
+"""
 function unpack(packed_array::AbstractArray{<:Any,N}, packed_shapes, pattern::PackPattern{N}) where N
     checkpacking(pattern)
     inds = Iterators.accumulate(+, Iterators.map(prod, packed_shapes))
     unpacked_arrays = map(Iterators.flatten((0, inds)), inds, packed_shapes) do i, j, ps
-        A = selectdim(packed_array, find_wildcard(pattern), i+1:j)
-        reshape(A, splice(size(A), find_wildcard(pattern), ps))
+        unpacked_array = selectdim(packed_array, find_wildcard(pattern), i+1:j)
+        return collect(reshape(unpacked_array, splice(size(unpacked_array), find_wildcard(pattern), ps)))
     end
     return unpacked_arrays
 end

src/rearrange.jl

@@ -55,21 +55,20 @@ 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
 (2, 7, 15)
 ```
 """
-function rearrange(x, @nospecialize pattern::Pattern; context...)
-    left, right = pattern
+function rearrange(x, (left, right); context...)
     (!isempty(extract(typeof(..), left)) || !isempty(extract(typeof(..), right))) && throw(ArgumentError("Ellipses (..) are currently not supported"))
     left_names, right_names = extract(Symbol, left), extract(Symbol, right)
     reshaped_in = reshape_in(x, left; context...)

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/