Repository: JuliaML/MLUtils.jl (julia)

Issue:
extend batch to handle vectors of tuples, named tuples, dicts

Pull request description:
Fix #23

Patch:
diff --git a/src/observation.jl b/src/observation.jl
index 067cd68..d686012 100644
--- a/src/observation.jl
+++ b/src/observation.jl
@@ -13,13 +13,15 @@ function numobs end
 Return the observations corresponding to the observation-index `idx`.
 Note that `idx` can be any type as long as `data` has defined
 `getobs` for that type.
+
 The returned observation(s) should be in the form intended to
 be passed as-is to some learning algorithm. There is no strict
 interface requirement on how this "actual data" must look like.
+
 Every author behind some custom data container can make this
 decision themselves.
-The output should be consistent when `idx` is a scalar vs vector.
 
+The output should be consistent when `idx` is a scalar vs vector.
 
 See also [`getobs!`](@ref) and [`numobs`](@ref) 
 """
@@ -35,9 +37,10 @@ getobs(data) = data
 Inplace version of `getobs(data, idx)`. If this method
 is defined for the type of `data`, then `buffer` should be used
 to store the result, instead of allocating a dedicated object.
+
 Implementing this function is optional. In the case no such
 method is provided for the type of `data`, then `buffer` will be
-*ignored* and the result of `getobs` returned. This could be
+*ignored* and the result of [`getobs`](@ref) returned. This could be
 because the type of `data` may not lend itself to the concept
 of `copy!`. Thus, supporting a custom `getobs!` is optional
 and not required.
diff --git a/src/utils.jl b/src/utils.jl
index ecd5ee1..f0a7a8d 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -147,18 +147,32 @@ end
 """
     batch(xs)
 
-Batch the arrays in `xs` into a single array.
-See also [`unbatch`](@ref)
+Batch the arrays in `xs` into a single array with 
+an extra dimension.
+
+If the elements of `xs` are tuples, named tuples, or dicts, 
+the output will be of the same type. 
+
+See also [`unbatch`](@ref).
+
 # Examples
+
 ```jldoctest
-julia> batch([[1,2,3],[4,5,6]])
+julia> batch([[1,2,3], 
+              [4,5,6]])
 3×2 Matrix{Int64}:
  1  4
  2  5
  3  6
+
+julia> batch([(a=[1,2], b=[3,4])
+               (a=[5,6], b=[7,8])]) 
+(a = [1 5; 2 6], b = [3 7; 4 8])
 ```
 """
 function batch(xs)
+# Fallback for generric iterables
+    @assert length(xs) > 0 "Input should be non-empty" 
     data = first(xs) isa AbstractArray ?
         similar(first(xs), size(first(xs))..., length(xs)) :
         Vector{eltype(xs)}(undef, length(xs))
@@ -170,6 +184,32 @@ end
 
 batchindex(xs, i) = (reverse(Base.tail(reverse(axes(xs))))..., i)
 
+function batch(xs::AbstractVector{<:AbstractArray{T,N}}) where {T, N}
+    return stack(xs, N+1)
+end
+
+function batch(xs::Vector{<:Tuple})
+    @assert length(xs) > 0 "Input should be non-empty"
+    n = length(first(xs))
+    @assert all(length.(xs) .== n) "Cannot batch tuples with different lengths"
+    return ntuple(i -> batch([x[i] for x in xs]), n)
+end
+
+function batch(xs::Vector{<:NamedTuple})
+    @assert length(xs) > 0 "Input should be non-empty"
+    all_keys = [sort(collect(keys(x))) for x in xs]
+    ks = all_keys[1]
+    @assert all(==(ks), all_keys) "Cannot batch named tuples with different keys"
+    NamedTuple(k => batch([x[k] for x in xs]) for k in ks)
+end
+
+function batch(xs::Vector{<:Dict})
+    @assert length(xs) > 0 "Input should be non-empty"
+    all_keys = [sort(collect(keys(x))) for x in xs]
+    ks = all_keys[1]
+    @assert all(==(ks), all_keys) "cannot batch dicts with different keys"
+    Dict(k => batch([x[k] for x in xs]) for k in ks)
+end
 
 """
     unbatch(x)


Test patch:
diff --git a/test/utils.jl b/test/utils.jl
index a018d72..79f8a5a 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -1,41 +1,69 @@
 @testset "unsqueeze" begin
-  x = randn(2, 3, 2)
-  @test @inferred(unsqueeze(x, 1)) == reshape(x, 1, 2, 3, 2)
-  @test @inferred(unsqueeze(x, 2)) == reshape(x, 2, 1, 3, 2)
-  @test @inferred(unsqueeze(x, 3)) == reshape(x, 2, 3, 1, 2)
-  @test @inferred(unsqueeze(x, 4)) == reshape(x, 2, 3, 2, 1)
+    x = randn(2, 3, 2)
+    @test @inferred(unsqueeze(x, 1)) == reshape(x, 1, 2, 3, 2)
+    @test @inferred(unsqueeze(x, 2)) == reshape(x, 2, 1, 3, 2)
+    @test @inferred(unsqueeze(x, 3)) == reshape(x, 2, 3, 1, 2)
+    @test @inferred(unsqueeze(x, 4)) == reshape(x, 2, 3, 2, 1)
 
-  @test unsqueeze(2)(x) == unsqueeze(x, 2)
+    @test unsqueeze(2)(x) == unsqueeze(x, 2)
 end
 
 @testset "stack and unstack" begin
-  x = randn(3,3)
-  stacked = stack([x, x], 2)
-  @test size(stacked) == (3,2,3)
-
-  stacked_array=[ 8 9 3 5; 9 6 6 9; 9 1 7 2; 7 4 10 6 ]
-  unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
-  @test unstack(stacked_array, 2) == unstacked_array
-  @test stack(unstacked_array, 2) == stacked_array
-  @test stack(unstack(stacked_array, 1), 1) == stacked_array
+    x = randn(3,3)
+    stacked = stack([x, x], 2)
+    @test size(stacked) == (3,2,3)
+
+    stacked_array=[ 8 9 3 5; 9 6 6 9; 9 1 7 2; 7 4 10 6 ]
+    unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
+    @test unstack(stacked_array, 2) == unstacked_array
+    @test stack(unstacked_array, 2) == stacked_array
+    @test stack(unstack(stacked_array, 1), 1) == stacked_array
 end
 
 @testset "batch and unbatch" begin
-  stacked_array=[ 8 9 3 5
-                  9 6 6 9
-                  9 1 7 2
-                  7 4 10 6 ]
-  unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
-  @test unbatch(stacked_array) == unstacked_array
-  @test batch(unstacked_array) == stacked_array
-
-  # no-op for vector of non-arrays
-  @test batch([1,2,3]) == [1,2,3]
-  @test unbatch([1,2,3]) == [1,2,3]
-
-  # generic iterable
-  @test batch(ones(2) for i=1:3) == ones(2, 3)
-  @test unbatch(ones(2, 3)) == [ones(2) for i=1:3]
+    stacked_array=[ 8 9 3 5
+                    9 6 6 9
+                    9 1 7 2
+                    7 4 10 6 ]
+    unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
+    @test unbatch(stacked_array) == unstacked_array
+    @test batch(unstacked_array) == stacked_array
+
+    # no-op for vector of non-arrays
+    @test batch([1,2,3]) == [1,2,3]
+    @test unbatch([1,2,3]) == [1,2,3]
+
+    # generic iterable
+    @test batch(ones(2) for i=1:3) == ones(2, 3)
+    @test unbatch(ones(2, 3)) == [ones(2) for i=1:3]
+
+    @testset "tuple" begin
+        @test batch([(1,2), (3,4)]) == ([1,3], [2,4])
+        @test batch([([1,2], [3,4]), ([5,6],[7,8])]) == ([1 5
+                                                          2 6], 
+                                                         [3 7
+                                                          4 8])      
+    end
+
+    @testset "named tuple" begin
+        @test batch([(a=1,b=2), (a=3,b=4)]) == (a=[1,3], b=[2,4])
+        @test batch([(a=1,b=2), (b=4,a=3)]) == (a=[1,3], b=[2,4])
+        nt = [(a=[1,2], b=[3,4]), (a=[5,6],b=[7,8])]
+        @test batch(nt) == (a = [1 5
+                                2 6], 
+                            b = [3 7
+                                4 8])      
+    end
+
+    @testset "dict" begin
+        @test batch([Dict(:a=>1,:b=>2), Dict(:a=>3,:b=>4)]) == Dict(:a=>[1,3], :b=>[2,4])
+        @test batch([Dict(:a=>1,:b=>2), Dict(:b=>4,:a=>3)]) == Dict(:a=>[1,3], :b=>[2,4])
+        d = [Dict(:a=>[1,2], :b=>[3,4]), Dict(:a=>[5,6],:b=>[7,8])]
+        @test batch(d) == Dict(:a => [1 5
+                                      2 6], 
+                                :b => [3 7
+                                      4 8])    
+    end
 end
 
 @testset "flatten" begin