Merge #1726

1726: add unbatch r=CarloLucibello a=CarloLucibello

Add `unbatch` for convenience and consistency (we currently have `stack/unstack` but only `batch` with no `unbatch`). Can be specialized by downstream packages on their custom batched types. 

### PR Checklist

- [x] Tests are added
- [ ] Entry in NEWS.md
- [x] Documentation, if applicable
- [ ] API changes require approval from a committer (different from the author, if applicable)


Co-authored-by: Carlo Lucibello <carlo.lucibello@gmail.com>

Author: bors[bot]

docs/src/utilities.md

@@ -14,6 +14,7 @@ Flux.unstack
 Flux.chunk
 Flux.frequencies
 Flux.batch
+Flux.unbatch
 Flux.batchseq
 Base.rpad(v::AbstractVector, n::Integer, p)
 ```

src/utils.jl

@@ -542,6 +542,8 @@ squeezebatch(x) = reshape(x, head(size(x)))
 
 Batch the arrays in `xs` into a single array.
 
+See also [`unbatch`](@ref)
+
 # Examples
 ```jldoctest
 julia> Flux.batch([[1,2,3],[4,5,6]])
@@ -561,6 +563,28 @@ function batch(xs)
   return data
 end
 
+"""
+  unbatch(x)
+
+Reverse of the [`batch`](@ref) operation,
+unstacking the last dimension of the array `x`.
+
+See also [`unstack`](@ref).
+
+# Examples
+
+```jldoctest
+julia> Flux.unbatch([1 3 5 7; 
+                     2 4 6 8])
+4-element Vector{Vector{Int64}}:
+ [1, 2]
+ [3, 4]
+ [5, 6]
+ [7, 8]
+"""
+unbatch(x::AbstractArray) = unstack(x, ndims(x))
+unbatch(x::AbstractVector) = x
+
 """
 Return the given sequence padded with `p` up to a maximum length of `n`.
 

test/utils.jl

@@ -1,5 +1,7 @@
 using Flux
-using Flux: throttle, nfan, glorot_uniform, glorot_normal, kaiming_normal, kaiming_uniform, orthogonal, sparse_init, stack, unstack, Zeros
+using Flux: throttle, nfan, glorot_uniform, glorot_normal,
+             kaiming_normal, kaiming_uniform, orthogonal, 
+             sparse_init, stack, unstack, Zeros, batch, unbatch
 using StatsBase: var, std
 using Random
 using Test
@@ -329,6 +331,25 @@ end
   @test stack(unstack(stacked_array, 1), 1) == stacked_array
 end
 
+
+@testset "Batching" 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]
+end
+
 @testset "Param remapping" begin
   ls(dims...) = reshape(collect(Float32, 1:prod(dims)), dims...) # accepts dims in reverse order to Dense
   dl(nin, nout, bias) = Dense(ls(nout, nin), bias(nout))