generalize batchseq

Author: CarloLucibello

Project.toml

@@ -9,6 +9,7 @@ DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 FoldsThreads = "9c68100b-dfe1-47cf-94c8-95104e173443"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
@@ -23,8 +24,9 @@ DataAPI = "1.0"
 DelimitedFiles = "1.0"
 FLoops = "0.2"
 FoldsThreads = "0.1"
-SimpleTraits = "0.9"
+NNlib = "0.8"
 ShowCases = "0.1"
+SimpleTraits = "0.9"
 StatsBase = "0.33"
 Tables = "1.10"
 Transducers = "0.4"

docs/src/api.md

@@ -36,9 +36,7 @@ obsview
 ObsView
 ones_like
 oversample
-MLUtils.rpad
 randobs
-rpad(::AbstractVector, ::Integer, ::Any)
 shuffleobs
 splitobs
 stack

src/MLUtils.jl

@@ -17,6 +17,7 @@ using ChainRulesCore: @non_differentiable, unthunk, AbstractZero,
                       NoTangent, ZeroTangent, ProjectTo
 
 using SimpleTraits
+using NNlib
 
 @traitdef IsTable{X}
 @traitimpl IsTable{X} <- Tables.istable(X)
@@ -78,7 +79,6 @@ export batch,
        unsqueeze,
        unstack,
        zeros_like
-       # rpad
 
 include("Datasets/Datasets.jl")
 using .Datasets

src/deprecations.jl

@@ -1,4 +1,4 @@
-# Deprecations v0.1
+# Deprecated in v0.2
 @deprecate stack(x, dims) stack(x; dims=dims)
 @deprecate unstack(x, dims) unstack(x; dims=dims)
 @deprecate unsqueeze(x::AbstractArray, dims::Int) unsqueeze(x; dims=dims)
@@ -7,3 +7,10 @@
 @deprecate frequencies(x) group_counts(x)
 @deprecate eachbatch(data, batchsize; kws...) eachobs(data; batchsize, kws...)
 @deprecate eachbatch(data; size=1, kws...) eachobs(data; batchsize=size, kws...)
+
+# Deprecated in v0.3
+
+function Base.rpad(v::AbstractVector, n::Integer, p)
+    @warn "rpad is deprecated, NNlib.pad_zeros or NNlib.pad_constant instead"
+    return [v; fill(p, max(n - length(v), 0))]
+end

src/utils.jl

@@ -388,36 +388,10 @@ unbatch(x::AbstractArray) = [getobs(x, i) for i in 1:numobs(x)]
 unbatch(x::AbstractVector) = x
 
 """
-    rpad(v::AbstractVector, n::Integer, p)
-
-Return the given sequence padded with `p` up to a maximum length of `n`.
-
-# Examples
-
-```jldoctest
-julia> rpad([1, 2], 4, 0)
-4-element Vector{Int64}:
- 1
- 2
- 0
- 0
-
-julia> rpad([1, 2, 3], 2, 0)
-3-element Vector{Int64}:
- 1
- 2
- 3
-```
-"""
-Base.rpad(v::AbstractVector, n::Integer, p) = [v; fill(p, max(n - length(v), 0))]
-# TODO Piracy
-
-
-"""
-    batchseq(seqs, pad)
+    batchseq(seqs, val = 0)
 
 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`.
+item is a batch of `N`. Short sequences will be padded by `val`.
 
 # Examples
 
@@ -429,11 +403,24 @@ julia> batchseq([[1, 2, 3], [4, 5]], 0)
  [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]
+function batchseq(xs, val = 0, n = nothing)
+    n = n === nothing ? maximum(x -> size(x, ndims(x)), xs) : n
+    xs_ = [rpad_constant(x, n, val; dims=ndims(x)) for x in xs]
+    [batch([obsview(xs_[j], i) for j = 1:length(xs_)]) for i = 1:n]
 end
 
+function rpad_constant(x, n, val = 0; dims=:)
+    ns = Int[]
+    _dims = dims === Colon() ? (1:ndims(x)) : dims
+    _n = n isa Integer ? ntuple(i -> n, length(dims)) : n
+    for i in length(_n)
+        push!(ns, 0)
+        push!(ns, n - size(x, _dims[i]))
+    end
+    return pad_constant(x, tuple(ns...), val; dims)
+end
+
+
 """
     flatten(x::AbstractArray)
 

test/utils.jl

@@ -129,11 +129,6 @@ end
     @test d == Dict('a' => 1, 'b' => 2)
 end
 
-@testset "rpad" begin
-    @test rpad([1, 2], 4, 0) == [1, 2, 0, 0]
-    @test rpad([1, 2, 3], 2, 0) == [1,2,3]
-end
-
 @testset "batchseq" begin
     bs = batchseq([[1, 2, 3], [4, 5]], 0)
     @test bs[1] == [1, 4]
@@ -144,6 +139,11 @@ end
     @test bs[1] == [1, 4]
     @test bs[2] == [2, 5]
     @test bs[3] == [3, -1]
+
+    batchseq([ones(2,4), zeros(2, 3), ones(2,2)]) ==[[1.0 0.0 1.0; 1.0 0.0 1.0]
+                                                    [1.0 0.0 1.0; 1.0 0.0 1.0]
+                                                    [1.0 0.0 0.0; 1.0 0.0 0.0]
+                                                    [1.0 0.0 0.0; 1.0 0.0 0.0]]
 end
 
 @testset "ones_like" begin