Merge pull request #106 from JuliaML/cl/partial

fix partial=false bug

Author: CarloLucibello

.gitignore

@@ -3,3 +3,4 @@
 *.jl.mem
 /Manifest.toml
 docs/build
+.vscode

Project.toml

@@ -1,7 +1,7 @@
 name = "MLUtils"
 uuid = "f1d291b0-491e-4a28-83b9-f70985020b54"
 authors = ["Carlo Lucibello <[email protected]> and contributors"]
-version = "0.2.7"
+version = "0.2.8"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/batchview.jl

@@ -87,10 +87,8 @@ struct BatchView{TElem,TData,TCollate} <: AbstractDataContainer
     batchsize::Int
     count::Int
     partial::Bool
-    imax::Int
 end
 
-
 function BatchView(data::T; batchsize::Int=1, partial::Bool=true, collate=Val(nothing)) where {T}
     n = numobs(data)
     if n < batchsize
@@ -102,9 +100,8 @@ function BatchView(data::T; batchsize::Int=1, partial::Bool=true, collate=Val(no
         throw(ArgumentError("`collate` must be one of `nothing`, `true` or `false`."))
     end
     E = _batchviewelemtype(data, collate)
-    imax = partial ? n : n - batchsize + 1
     count = partial ? ceil(Int, n / batchsize) : floor(Int, n / batchsize)
-    BatchView{E,T,typeof(collate)}(data, batchsize, count, partial, imax)
+    BatchView{E,T,typeof(collate)}(data, batchsize, count, partial)
 end
 
 _batchviewelemtype(::TData, ::Val{nothing}) where TData =
@@ -155,10 +152,10 @@ Base.iterate(A::BatchView, state = 1) =
     (state > numobs(A)) ? nothing : (A[state], state + 1)
 
 # Helper function to translate a batch-index into a range of observations.
-@inline function _batchrange(a::BatchView, batchindex::Int)
-    @boundscheck (batchindex > a.count || batchindex < 0) && throw(BoundsError())
-    startidx = (batchindex - 1) * a.batchsize + 1
-    endidx = min(a.imax, startidx + a.batchsize -1)
+@inline function _batchrange(A::BatchView, batchindex::Int)
+    @boundscheck (batchindex > A.count || batchindex < 0) && throw(BoundsError())
+    startidx = (batchindex - 1) * A.batchsize + 1
+    endidx = min(numobs(parent(A)), startidx + A.batchsize -1)
     return startidx:endidx
 end
 

test/batchview.jl

@@ -107,4 +107,13 @@ using MLUtils: obsview
     #         @test eltype(@inferred(BatchView(ObsView(var)))[1]) <: SubArray
     #     end
     # end
+
+    @testset "partial=false" begin
+        x = [1:12;]
+        bv = BatchView(x, batchsize=5, partial=false)
+        @test length(bv) == 2
+        @test bv[1] == 1:5
+        @test bv[2] == 6:10
+        @test_throws BoundsError bv[3]
+    end
 end

test/dataloader.jl

@@ -81,6 +81,14 @@
         end
     end
 
+    @testset "partial=false" begin
+        x = [1:12;]
+        d = DataLoader(x, batchsize=5, partial=false) |> collect
+        @test length(d) == 2
+        @test d[1] == 1:5
+        @test d[2] == 6:10
+    end
+
     @testset "shuffle & rng" begin
         X4 = rand(2, 1000)
         d1 = DataLoader(X4, batchsize=2; shuffle=true)
@@ -97,6 +105,7 @@
         @test first(d1) == first(d2)
     end
 
+
     # numobs/getobs compatibility
     d = DataLoader(CustomType(), batchsize=2)
     @test first(d) == [1, 2]