multi-array input

Author: AntonOresten

src/rearrange.jl

@@ -45,8 +45,11 @@ end
 
 reshape_out(x, ::Tuple{Vararg{Symbol}}) = x
 
+# TODO: support taking vectors of arrays
+
 """
-    rearrange(x::AbstractArray, left --> right)
+    rearrange(array::AbstractArray, left --> right; context...)
+    rearrange(arrays, left --> right; context...)
 
 Rearrange the axes of `x` according to the pattern specified by `left --> right`.
 
@@ -68,11 +71,17 @@ julia> reshape(permutedims(reshape(rand(2,3,35), 2,3,5,7), (1,4,3,2)), 2,7,5*3)
 (2, 7, 15)
 ```
 """
-function rearrange(x, (left, right); context...)
+function rearrange(x::AbstractArray, pattern::Pattern; context...)
+    @nospecialize pattern
+    left, right = pattern
     (!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...)
     permuted = permute(reshaped_in, left_names, right_names)
     reshaped_out = reshape_out(permuted, right)
     return reshaped_out
 end
+
+rearrange(x::AbstractArray{<:AbstractArray}, pattern::Pattern; context...) = rearrange(stack(x), pattern; context...)
+
+rearrange(x, pattern::Pattern; context...) = rearrange(stack(x), pattern; context...)

test/runtests.jl

@@ -74,6 +74,11 @@ using Test, Statistics
         @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))
 
+        x = rand(2,3,5)
+        @test rearrange([x, x], (:a, :b, :c, :d) --> (:c, :b, :a, :d)) == permutedims(stack([x, x]), (3,2,1,4))
+        @test rearrange(reshape([x, x], 1, 2), (:a, :b, :c, 1, :d) --> (:c, :b, :a, :d)) == permutedims(reshape(cat(x, x, dims=5), 2,3,5,2), (3,2,1,4))
+        @test rearrange((x, x), (:a, :b, :c, :d) --> (:c, :b, :a, :d)) == permutedims(cat(x, x, dims=4), (3,2,1,4))
+
         @testset "Python API reference parity" begin
             # see https://einops.rocks/api/rearrange/