Add pack and unpack

Author: AntonOresten

src/Einops.jl

@@ -22,10 +22,13 @@ export @einops_str
 include("rearrange.jl")
 export rearrange
 
+include("pack_unpack.jl")
+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")
 
-# TODO: einsum, pack, unpack
+# TODO: einsum
 
 end

src/einops_str.jl

@@ -1,14 +1,13 @@
 function parse_pattern(pattern::AbstractString)
-    occursin("->", pattern) ||
-        throw(ArgumentError("pattern must contain \"->\" (got \"$pattern\")"))
+    occursin("->", pattern) || return tokenize_packing_pattern(pattern)
 
     lhs, rhs = strip.(split(pattern, "->"; limit = 2))
-    lhs_axes = tokenise_side(lhs)
-    rhs_axes = tokenise_side(rhs)
+    lhs_axes = tokenize_side(lhs)
+    rhs_axes = tokenize_side(rhs)
     return Tuple(lhs_axes) --> Tuple(rhs_axes)
 end
 
-function tokenise_side(side::AbstractString)
+function tokenize_side(side::AbstractString)
     tokens = Any[]
     stack  = Vector{Any}[]
     buf    = IOBuffer()
@@ -73,8 +72,13 @@ function tokenise_side(side::AbstractString)
     return Tuple(tokens)
 end
 
+function tokenize_packing_pattern(pattern::AbstractString)
+    Tuple(map(s -> s == "*" ? (*) : Symbol(s), filter(!isempty, split(pattern, ' '))))
+end
+
 """
-    einops"... -> ..."
+    einops"a (b c) -> (c b a)"
+    einops"i j * k"
 
 For parity with Python implementation.
 
@@ -86,6 +90,9 @@ julia> einops"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)
+
+julia> einops"i j * k" # for packing
+(:i, :j, *, :k)
 ```
 """
 macro einops_str(pattern)

src/pack_unpack.jl

@@ -0,0 +1,31 @@
+const Wildcard = typeof(*)
+const PackPattern{N} = NTuple{N,Union{Symbol,Wildcard}}
+
+checkpacking(pattern::PackPattern) = count(x -> x isa Wildcard, pattern) == 1 || error("Only one wildcard (*) is allowed in the pattern")
+find_wildcard(pattern::PackPattern) = findfirst(x -> x isa Wildcard, pattern)
+
+size_before_wildcard(dims::Dims, pattern::PackPattern) = dims[1:find_wildcard(pattern)-1]
+size_after_wildcard(dims::Dims, pattern::PackPattern) = dims[end-length(pattern)+find_wildcard(pattern)+1:end]
+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
+    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]
+    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?
+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))
+    end
+    return unpacked_arrays
+end

test/runtests.jl

@@ -16,18 +16,28 @@ using Test, Statistics
     end
 
     @testset "einops string tokenization" 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_throws "'.'" Einops.parse_pattern("-> .")
-        @test_throws "'('" Einops.parse_pattern("-> (")
-        @test_throws "')'" Einops.parse_pattern("-> )")
-        @test_throws "->" Einops.parse_pattern("")
+
+        @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_throws "'.'" Einops.parse_pattern("-> .")
+            @test_throws "'('" Einops.parse_pattern("-> (")
+            @test_throws "')'" Einops.parse_pattern("-> )")
+        end
+
+        @testset "packing pattern" begin
+            @test einops"i j * k" == (:i, :j, *, :k)
+            @test einops" i  j  *  k " == (:i, :j, *, :k)
+            @test einops"* i" == (*, :i)
+            @test einops"i *" == (:i, *)
+        end
+
     end
 
     @testset "rearrange" begin
@@ -172,4 +182,25 @@ using Test, Statistics
 
     end
 
+    @testset "pack_unpack" begin
+
+        @testset "Python API reference parity" begin
+            # see https://einops.rocks/api/pack_unpack/
+
+            inputs = [rand(2, 3, 5), rand(2, 3, 7, 5), rand(2, 3, 7, 9, 5)]
+            @test begin
+                packed, ps = pack(inputs, einops"i j * k")
+                packed |> size == (2, 3, 71, 5) && ps == [(), (7,), (7, 9)]
+            end
+
+            @test begin
+                packed, ps = pack(inputs, einops"i j * k")
+                inputs_unpacked = unpack(packed, ps, einops"i j * k")
+                all(inputs .== inputs_unpacked)
+            end
+
+        end
+
+    end
+
 end