Make DataLoader wrap data with ObsView (#73)

* fix obsview for cuarrays

* more tests

* obsview abstractrange

* view also for SubArray

* fix some tests

* just fix in DataLoader

* tests

* fix

Project.toml

@@ -1,7 +1,7 @@
 name = "MLUtils"
 uuid = "f1d291b0-491e-4a28-83b9-f70985020b54"
 authors = ["Carlo Lucibello <carlo.lucibello@gmail.com> and contributors"]
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/dataloader.jl

@@ -83,7 +83,9 @@ function DataLoader(data; batchsize=1, shuffle=false, partial=true, rng=GLOBAL_R
 end
 
 @propagate_inbounds function Base.iterate(d::DataLoader)
-    data = d.data
+    # Wrapping with ObsView in order to work around
+    # issue https://github.com/FluxML/Flux.jl/issues/1935
+    data = ObsView(d.data)
     if d.shuffle
         data = shuffleobs(d.rng, data)
     end

test/dataloader.jl

@@ -128,4 +128,19 @@
         c = collect(dloader)
         @test eltype(c) == Vector{Int}
     end
+
+    # https://github.com/FluxML/Flux.jl/issues/1935
+    @testset "no views of arrays" begin
+        x = CustomArrayNoView(6)
+        @test_throws ErrorException view(x, 1:2)
+
+        d = DataLoader(x)
+        @test length(collect(d)) == 6 # succesfull iteration
+
+        d = DataLoader(x, batchsize=2, shuffle=false)
+        @test length(collect(d)) == 3 # succesfull iteration
+
+        d = DataLoader(x, batchsize=2, shuffle=true)
+        @test length(collect(d)) == 3 # succesfull iteration
+    end
 end

test/obsview.jl

@@ -50,6 +50,7 @@
         end
     end
 
+
     @testset "custom types" begin
         @test_throws MethodError ObsView(EmptyType())
         @test_throws MethodError ObsView(EmptyType(), 1:10)
@@ -60,6 +61,9 @@
         @test getobs(ObsView(CustomType())) == collect(1:15)
         @test getobs(ObsView(CustomType(),1:10),10) == 10
         @test getobs(ObsView(CustomType(),1:10),[3,5]) == [3,5]
+
+        @test obsview(CustomArray(5)) isa SubArray
+        @test getobs(obsview(CustomArray(5)), 1:2) == CustomArray(2) 
     end
 end
 

test/runtests.jl

@@ -42,6 +42,37 @@ MLUtils.numobs(::CustomType) = 15
 MLUtils.getobs(::CustomType, i::Int) = i
 MLUtils.getobs(::CustomType, i::AbstractVector) = collect(i)
 
+struct CustomArray{T,N} <: AbstractArray{T,N} 
+    length::Int
+    size::NTuple{N, Int}
+
+    function CustomArray{T}(n...) where T
+        N = length(n)
+        new{T,N}(prod(n), Tuple(n))
+    end
+end
+CustomArray(n...) = CustomArray{Float32}(n...)
+Base.length(x::CustomArray) = x.length
+Base.size(x::CustomArray) = x.size
+Base.getindex(x::CustomArray{T}, i::Int) where T = T(1)
+Base.getindex(x::CustomArray{T, 1}, i::AbstractVector{Int}) where T = CustomArray{T}(length(i))
+
+struct CustomArrayNoView{T,N} <: AbstractArray{T,N} 
+    length::Int
+    size::NTuple{N, Int}
+
+    function CustomArrayNoView{T}(n...) where T
+        N = length(n)
+        new{T,N}(prod(n), Tuple(n))
+    end
+end
+CustomArrayNoView(n...) = CustomArrayNoView{Float32}(n...)
+Base.length(x::CustomArrayNoView) = x.length
+Base.size(x::CustomArrayNoView) = x.size
+Base.getindex(x::CustomArrayNoView{T}, i::Int) where T = T(1)
+Base.getindex(x::CustomArrayNoView{T, 1}, i::AbstractVector{Int}) where T = CustomArrayNoView{T}(length(i))
+Base.view(x::CustomArrayNoView, i...) = error("view not allowed")
+
 # --------------------------------------------------------------------
 
 include("test_utils.jl")