batching docs

This commit is contained in:
Mike J Innes 2017-10-18 16:21:15 +01:00
parent b817ce632c
commit e82428bb83

View File

@ -11,8 +11,23 @@ unstack(xs, dim) = [slicedim(xs, dim, i) for i = 1:size(xs, dim)]
batchindex(xs, i) = (reverse(Base.tail(reverse(indices(xs))))..., i)
"""
batch(xs)
Batch the arrays in `xs` into a single array.
```julia
julia> batch([[1,2,3],[4,5,6]])
3×2 Array{Int64,2}:
1 4
2 5
3 6
```
"""
function batch(xs)
data = similar(first(xs), size(first(xs))..., length(xs))
data = first(xs) isa AbstractArray ?
similar(first(xs), size(first(xs))..., length(xs)) :
Vector{eltype(xs)}(length(xs))
for (i, x) in enumerate(xs)
data[batchindex(data, i)...] = x
end
@ -21,7 +36,21 @@ end
Base.rpad(v::AbstractVector, n::Integer, p) = [v; fill(p, max(n - length(v), 0))]
function batchseq(xs, pad, n = maximum(length(x) for x in xs))
"""
batchseq(seqs, pad)
Take a list of `N` sequences, and turn them into a single sequence where each
item is a batch of `N`. Short sequences will be padded by `pad`.
```julia
julia> batchseq([[1, 2, 3], [4, 5]], 0)
3-element Array{Array{Int64,1},1}:
[1, 4]
[2, 5]
[3, 0]
```
"""
function batchseq(xs, pad = nothing, n = maximum(length(x) for x in xs))
xs_ = [rpad(x, n, pad) for x in xs]
[batch([xs_[j][i] for j = 1:length(xs_)]) for i = 1:n]
end