Repository: JuliaML/MLUtils.jl
Pull request #34 (2022-02-03 09:03:11)

remove eachbatch and rebase DataLoader on eachobs
With this, we still have a Flux compliant `DataLoader` but now the implementation is based on `eachobs`.

Also removing `eachbatch` in favor of   `eachobs(data; batchsize)`

Related to #33 #30 




Patch:
diff --git a/src/MLUtils.jl b/src/MLUtils.jl
index 9495ddc..6eff7ec 100644
--- a/src/MLUtils.jl
+++ b/src/MLUtils.jl
@@ -22,9 +22,8 @@ include("batchview.jl")
 export batchsize,
        batchview, BatchView
 
-include("dataiterator.jl")
-export eachobs, 
-       eachbatch
+include("eachobs.jl")
+export eachobs
 
 include("dataloader.jl")
 export DataLoader
diff --git a/src/batchview.jl b/src/batchview.jl
index efa1014..435592c 100644
--- a/src/batchview.jl
+++ b/src/batchview.jl
@@ -1,6 +1,6 @@
 """
-    BatchView(data, size; partial=true)
-    BatchView(data; size=1, partial=true)
+    BatchView(data, batchsize; partial=true)
+    BatchView(data; batchsize=1, partial=true)
 
 
 Description
@@ -9,11 +9,11 @@ Description
 Create a view of the given `data` that represents it as a vector
 of batches. Each batch will contain an equal amount of
 observations in them. The batch-size
-can be specified using the  parameter `size`.
+can be specified using the  parameter `batchsize`.
 In the case that the size of the dataset is not dividable by the
-specified `size`, the remaining observations will
+specified `batchsize`, the remaining observations will
 be ignored if `partial=false`. If  `partial=true` instead 
-the last batch size can be slight smaller.
+the last batch-size can be slight smaller.
 
 Note that any data access is delayed until `getindex` is called,
 and even `getindex` returns the result of [`obsview`](@ref)
@@ -30,11 +30,12 @@ Arguments
     type as long as it implements [`getobs`](@ref) and
     [`numobs`](@ref) (see Details for more information).
 
-- **`size`** : The batch-size of each batch. I.e. the number of
-    observations that each batch must contain.
+- **`batchsize`** : The batch-size of each batch. 
+    It is the number of observations that each batch must contain
+    (except possibly for the last one).
 
 - **`partial`** : If `partial=false` and the number of observations is 
-    not divisible by the batch size, then the last mini-batch is dropped.
+    not divisible by the batch-size, then the last mini-batch is dropped.
 
 Details
 ========
@@ -51,14 +52,14 @@ Examples
 using MLUtils
 X, Y = MLUtils.load_iris()
 
-A = batchview(X, size = 30)
+A = batchview(X, batchsize = 30)
 @assert typeof(A) <: BatchView <: AbstractVector
 @assert eltype(A) <: SubArray{Float64,2}
 @assert length(A) == 5 # Iris has 150 observations
 @assert size(A[1]) == (4,30) # Iris has 4 features
 
 # 5 batches of size 30 observations
-for x in batchview(X, size = 30)
+for x in batchview(X, batchsize = 30)
     @assert typeof(x) <: SubArray{Float64,2}
     @assert numobs(x) === 30
 end
@@ -67,7 +68,7 @@ end
 # Note that the iris dataset has 150 observations,
 # which means that with a batchsize of 20, the last
 # 10 observations will be ignored
-for (x,y) in batchview((X,Y), size = 20, partial=false)
+for (x,y) in batchview((X,Y), batchsize = 20, partial=false)
     @assert typeof(x) <: SubArray{Float64,2}
     @assert typeof(y) <: SubArray{String,1}
     @assert numobs(x) == numobs(y) == 20
@@ -75,7 +76,7 @@ end
 
 
 # randomly assign observations to one and only one batch.
-for (x,y) in batchview(shuffleobs((X,Y)), size=20)
+for (x,y) in batchview(shuffleobs((X,Y)), batchsize=20)
     @assert typeof(x) <: SubArray{Float64,2}
     @assert typeof(y) <: SubArray{String,1}
 end
@@ -89,7 +90,7 @@ see also
 """
 struct BatchView{TElem,TData} <: AbstractDataContainer
     data::TData
-    size::Int
+    batchsize::Int
     count::Int
     partial::Bool
     imax::Int
@@ -98,17 +99,16 @@ end
 @doc (@doc BatchView)
 const batchview = BatchView
 
-function BatchView(data::T; size::Int=1, partial::Bool=true) where {T}
-    # nsize, ncount = _compute_batch_settings(data, size, count, upto)
+function BatchView(data::T; batchsize::Int=1, partial::Bool=true) where {T}
     n = numobs(data)
-    if n < size
-        @warn "Number of observations less than batch size, decreasing the batch size to $n"
-        size = n
+    if n < batchsize
+        @warn "Number of observations less than batch-size, decreasing the batch-size to $n"
+        batchsize = n
     end
-    E = typeof(datasubset(data, 1:size))
-    imax = partial ? n : n - size + 1
-    count = partial ? ceil(Int, n / size) : floor(Int, n / size)
-    BatchView{E,T}(data, size, count, partial, imax)
+    E = typeof(datasubset(data, 1:batchsize))
+    imax = partial ? n : n - batchsize + 1
+    count = partial ? ceil(Int, n / batchsize) : floor(Int, n / batchsize)
+    BatchView{E,T}(data, batchsize, count, partial, imax)
 end
 
 """
@@ -116,20 +116,24 @@ end
 
 Return the fixed size of each batch in `data`.
 """
-batchsize(A::BatchView) = A.size
+batchsize(A::BatchView) = A.batchsize
 
 numobs(A::BatchView) = A.count
 getobs(A::BatchView) = getobs(A.data)
-getobs(A::BatchView, i) = getobs(A[i])
+getobs(A::BatchView, i::Int) = getobs(A.data, _batchrange(A, i))
+
+function getobs(A::BatchView, is::AbstractVector)
+    obsindices = union((_batchrange(A, i) for i in is)...)::Vector{Int}
+    getobs(A.data, obsindices)
+end
 
 Base.parent(A::BatchView) = A.data
 Base.eltype(::BatchView{Tel}) where Tel = Tel
 
-Base.getindex(A::BatchView, batchindex::Int) =
-    datasubset(A.data, _batchrange(A, batchindex))
+Base.getindex(A::BatchView, i::Int) = datasubset(A.data, _batchrange(A, i))
 
-function Base.getindex(A::BatchView, batchindices::AbstractVector)
-    obsindices = union((_batchrange(A, bi) for bi in batchindices)...)::Vector{Int}
+function Base.getindex(A::BatchView, is::AbstractVector)
+    obsindices = union((_batchrange(A, i) for i in is)...)::Vector{Int}
     datasubset(A.data, obsindices)
 end
 
@@ -139,8 +143,8 @@ datasubset(A::BatchView, i) = A[i]
 # Helper function to translate a batch-index into a range of observations.
 function _batchrange(a::BatchView, batchindex::Int)
     (batchindex > a.count || batchindex < 0) && throw(BoundsError())
-    startidx = (batchindex - 1) * a.size + 1
-    endidx = min(numobs(a.data), startidx + a.size -1) 
+    startidx = (batchindex - 1) * a.batchsize + 1
+    endidx = min(numobs(a.data), startidx + a.batchsize -1) 
     return startidx:endidx
 end
 
@@ -148,7 +152,7 @@ function Base.showarg(io::IO, A::BatchView, toplevel)
     print(io, "batchview(")
     Base.showarg(io, parent(A), false)
     print(io, ", ")
-    print(io, "size=$(A.size), ")
+    print(io, "batchsize=$(A.batchsize), ")
     print(io, "partial=$(A.partial)")
     print(io, ')')
     toplevel && print(io, " with eltype ", nameof(eltype(A))) # simplify
diff --git a/src/dataiterator.jl b/src/dataiterator.jl
deleted file mode 100644
index 6319ad7..0000000
--- a/src/dataiterator.jl
+++ /dev/null
@@ -1,79 +0,0 @@
-"""
-    eachobs(data, buffer=false)
-
-Iterate over `data` one observation at a time. 
-
-If  `buffer=true` and supported by the type of `data`, 
-a buffer will be preallocated and reused for memory efficiency.
-A buffer object `b` can also be passed directly with `buffer=b`.
-
-```julia
-X = rand(4,100)
-for x in eachobs(X)
-    # loop entered 100 times
-    @assert typeof(x) <: Vector{Float64}
-    @assert size(x) == (4,)
-end
-```
-
-Multiple variables are supported (e.g. for labeled data)
-
-```julia
-for (x,y) in eachobs((X,Y))
-    # ...
-end
-```
-
-See [`eachbatch`](@ref) for a mini-batch version.
-"""
-function eachobs(data; buffer=false)
-    if buffer === false 
-        return (getobs(data, i) for i in 1:numobs(data))
-    else
-        if buffer === true && numobs(data) > 0
-            buffer = getobs(data, 1)
-        end
-        return (getobs!(buffer, data, i) for i in 1:numobs(data))
-    end
-end
-
-
-# --------------------------------------------------------------------
-
-"""
-    eachbatch(data; size=1, partial=true, buffer=false)
-    
-Iterate over `data` one batch at a time. If supported by the type
-of `data`, a buffer will be preallocated and reused for memory
-efficiency.
-
-
-The  batch-size can be either provided directly using
-`size`. 
-
-```julia
-X = rand(4,150)
-for x in eachbatch(X, size=10)
-    # loop entered 15 times
-    @assert typeof(x) <: Matrix{Float64}
-    @assert size(x) == (4,10)
-end
-```
-
-Multiple variables are supported (e.g. for labeled data):
-
-```julia
-for (x,y) in eachbatch((X,Y))
-    # ...
-end
-```
-
-See also [`batchview`](@ref) and [`eachobs`](@ref).
-"""
-function eachbatch(data; buffer=false, kws...)
-    batched = BatchView(data; kws...)
-    return eachobs(batched; buffer)
-end
-
-eachbatch(data, size; kws...) = eachbatch(data; size, kws...)
-
diff --git a/src/dataloader.jl b/src/dataloader.jl
index 0bfe8e3..bde1724 100644
--- a/src/dataloader.jl
+++ b/src/dataloader.jl
@@ -5,8 +5,6 @@ struct DataLoader{D,R<:AbstractRNG}
     batchsize::Int
     nobs::Int
     partial::Bool
-    imax::Int
-    indices::Vector{Int}
     shuffle::Bool
     rng::R
 end
@@ -76,38 +74,35 @@ julia> foreach(println∘summary, DataLoader(rand(Int8, 10, 64), batchsize=30))
 """
 function DataLoader(data; batchsize=1, shuffle=false, partial=true, rng=GLOBAL_RNG)
     batchsize > 0 || throw(ArgumentError("Need positive batchsize"))
-
-    n = numobs(data)
-    if n < batchsize
-        @warn "Number of observations less than batchsize, decreasing the batchsize to $n"
-        batchsize = n
+    nobs = numobs(data)
+    if nobs < batchsize
+        @warn "Number of observations less than batchsize, decreasing the batchsize to $nobs"
+        batchsize = nobs
     end
-    imax = partial ? n : n - batchsize + 1
-    DataLoader(data, batchsize, n, partial, imax, [1:n;], shuffle, rng)
+    DataLoader(data, batchsize, nobs, partial, shuffle, rng)
 end
 
-# returns data in d.indices[i+1:i+batchsize]
-@propagate_inbounds function Base.iterate(d::DataLoader, i=0)     
-    i >= d.imax && return nothing
-    if d.shuffle && i == 0
-        shuffle!(d.rng, d.indices)
+@propagate_inbounds function Base.iterate(d::DataLoader)
+    data = d.data
+    if d.shuffle
+        data = shuffleobs(d.rng, data)
     end
-    nexti = min(i + d.batchsize, d.nobs)
-    ids = d.indices[i+1:nexti]
-    batch = getobs(d.data, ids)
-    return (batch, nexti)
+    gen = eachobs(data; d.batchsize, d.partial)
+    res = iterate(gen)
+    res === nothing && return
+    res[1], (gen, res[2])
+end
+
+@propagate_inbounds function Base.iterate(d::DataLoader, state)     
+    gen, i = state
+    res = iterate(gen, i)
+    res === nothing && return
+    res[1], (gen, res[2])
 end
 
 function Base.length(d::DataLoader)
     n = d.nobs / d.batchsize
-    d.partial ? ceil(Int,n) : floor(Int,n)
+    d.partial ? ceil(Int, n) : floor(Int, n)
 end
 
-
-Base.eltype(::Type{<:DataLoader{D}}) where D = batchtype(D) 
-
-batchtype(D::Type) = Any
-batchtype(D::Type{<:AbstractArray}) = D
-batchtype(D::Type{<:Tuple})= Tuple{batchtype.(D.parameters)...}
-batchtype(D::Type{<:NamedTuple{K,V}}) where {K,V} = NamedTuple{K, batchtype(V)}
-batchtype(D::Type{<:Dict{K,V}}) where {K,V} = Dict{K, batchtype(V)}
+Base.IteratorEltype(d::DataLoader) = Base.EltypeUnknown()
\ No newline at end of file
diff --git a/src/deprecations.jl b/src/deprecations.jl
index f81e6f8..f9456e9 100644
--- a/src/deprecations.jl
+++ b/src/deprecations.jl
@@ -5,4 +5,5 @@
 @deprecate unsqueeze(dims::Int) unsqueeze(dims=dims)
 @deprecate labelmap(x) group_indices(x)
 @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...)
diff --git a/src/eachobs.jl b/src/eachobs.jl
new file mode 100644
index 0000000..bfc26a8
--- /dev/null
+++ b/src/eachobs.jl
@@ -0,0 +1,58 @@
+"""
+    eachobs(data, buffer=false, batchsize=-1, partial=true)
+
+Return an iterator over the observations in `data`. 
+
+# Arguments. 
+- `data`. The data to be iterated over. The data type has to implement 
+  [`numobs`](@ref) and [`getobs`](@ref). 
+- `buffer`. If `buffer=true` and supported by the type of `data`, 
+a buffer will be allocated and reused for memory efficiency.
+You can also pass a preallocated object to `buffer`.
+- `batchsize`. If less than 0, iterates over individual observation. 
+Otherwise, each iteration (except possibly the last) yields a mini-batch 
+containing `batchsize` observations.
+- `partial`. This argument is used only when `batchsize > 0`.
+  If `partial=false` and the number of observations is not divisible by the batchsize, 
+  then the last mini-batch is dropped.
+
+See also [`numobs`](@ref), [`getobs`](@ref) and [`batchview`](@ref).
+
+# Examples
+
+```julia
+X = rand(4,100)
+for x in eachobs(X)
+    # loop entered 100 times
+    @assert typeof(x) <: Vector{Float64}
+    @assert size(x) == (4,)
+end
+
+# mini-batch iterations
+for x in eachbatch(X, batchsize=10)
+    # loop entered 10 times
+    @assert typeof(x) <: Matrix{Float64}
+    @assert size(x) == (4,10)
+end
+
+# support for tuples, named tuples, dicts
+for (x, y) in eachobs((X, Y))
+    # ...
+end
+```
+"""
+function eachobs(data; buffer = false, batchsize::Int = -1, partial::Bool =true)
+    if batchsize > 0
+        data = BatchView(data; batchsize, partial)
+    end
+    if buffer === false 
+        gen = (getobs(data, i) for i in 1:numobs(data))
+    else
+        if buffer === true && numobs(data) > 0
+            buffer = getobs(data, 1)
+        end
+        gen = (getobs!(buffer, data, i) for i in 1:numobs(data))
+    end
+    return gen
+end
+


Test patch:
diff --git a/test/batchview.jl b/test/batchview.jl
index ef30424..ee83ec5 100644
--- a/test/batchview.jl
+++ b/test/batchview.jl
@@ -17,13 +17,13 @@ using MLUtils: datasubset
             @test_throws MethodError BatchView(var...)
             @test_throws MethodError BatchView(var, 16)
             
-            A = BatchView(var, size=3)
+            A = BatchView(var, batchsize=3)
             @test length(A) == 5
             @test batchsize(A) == 3
             @test numobs(A) == length(A)
             @test @inferred(parent(A)) === var
         end
-        A = BatchView(X, size=16)
+        A = BatchView(X, batchsize=16)
         @test length(A) == 1
         @test batchsize(A) == 15            
     end
@@ -32,24 +32,24 @@ using MLUtils: datasubset
     @testset "typestability" begin
         for var in (vars..., tuples..., Xs, ys)
             @test typeof(@inferred(BatchView(var))) <: BatchView
-            @test typeof(@inferred(BatchView(var, size=3))) <: BatchView
-            @test typeof(@inferred(BatchView(var, size=3, partial=true))) <: BatchView
-            @test typeof(@inferred(BatchView(var, size=3, partial=false))) <: BatchView
+            @test typeof(@inferred(BatchView(var, batchsize=3))) <: BatchView
+            @test typeof(@inferred(BatchView(var, batchsize=3, partial=true))) <: BatchView
+            @test typeof(@inferred(BatchView(var, batchsize=3, partial=false))) <: BatchView
         end
         @test typeof(@inferred(BatchView(CustomType()))) <: BatchView
     end
 
     @testset "AbstractArray interface" begin
         for var in (vars..., tuples..., Xs, ys)
-            A = BatchView(var, size=5)
+            A = BatchView(var, batchsize=5)
             @test_throws BoundsError A[-1]
             @test_throws BoundsError A[4]
             @test @inferred(numobs(A)) == 3
             @test @inferred(length(A)) == 3
             @test @inferred(batchsize(A)) == 5
             @test @inferred(size(A)) == (3,)
-            @test @inferred(getobs(A[2:3])) == getobs(BatchView(obsview(var, 6:15), size=5))
-            @test @inferred(getobs(A[[1,3]])) == getobs(BatchView(obsview(var, [1:5..., 11:15...]), size=5))
+            @test @inferred(getobs(A[2:3])) == getobs(BatchView(obsview(var, 6:15), batchsize=5))
+            @test @inferred(getobs(A[[1,3]])) == getobs(BatchView(obsview(var, [1:5..., 11:15...]), batchsize=5))
             @test @inferred(A[1]) == datasubset(var, 1:5)
             @test @inferred(A[2]) == datasubset(var, 6:10)
             @test @inferred(A[3]) == datasubset(var, 11:15)
@@ -60,7 +60,7 @@ using MLUtils: datasubset
             @test typeof(@inferred(collect(A))) <: Vector
         end
         for var in (vars..., tuples...)
-            A = BatchView(var, size=5)
+            A = BatchView(var, batchsize=5)
             @test @inferred(getobs(A)) == var
             @test A[2:3] == datasubset(var, [6:15;])
             @test A[[1,3]] == datasubset(var, [1:5..., 11:15...])
@@ -71,7 +71,7 @@ using MLUtils: datasubset
     @testset "subsetting" begin
         # for var in (vars..., tuples..., Xs, ys)
             for var in (vars[1], )
-            A = BatchView(var, size=3)
+            A = BatchView(var, batchsize=3)
             @test getobs(@inferred(obsview(A))) == @inferred(getobs(A))
             @test_throws BoundsError obsview(A,1:6)
             
@@ -82,12 +82,12 @@ using MLUtils: datasubset
             @test @inferred(size(S)) == (length(S),)
             @test getobs(@inferred(A[1:2])) == getobs(S)
             @test @inferred(getobs(A,1:2)) == getobs(S)
-            @test @inferred(getobs(S)) == getobs(BatchView(obsview(var,1:6),size=3))
+            @test @inferred(getobs(S)) == getobs(BatchView(obsview(var,1:6),batchsize=3))
             
             S = @inferred(ObsView(A, 1:2))
             @test typeof(S) <: ObsView
         end
-        A = BatchView(X, size=3)
+        A = BatchView(X, batchsize=3)
         @test typeof(A.data) <: Array
         S = @inferred(datasubset(A))
         S === A
diff --git a/test/dataiterator.jl b/test/dataiterator.jl
deleted file mode 100644
index 18dbd51..0000000
--- a/test/dataiterator.jl
+++ /dev/null
@@ -1,40 +0,0 @@
-@testset "eachobs" begin
-    for (i,x) in enumerate(eachobs(X))
-        @test x == X[:,i]
-    end
-
-    for (i,x) in enumerate(eachobs(X, buffer=true))
-        @test x == X[:,i]
-    end
-
-    b = zeros(size(X, 1))
-    for (i,x) in enumerate(eachobs(X, buffer=b))
-        @test x == X[:,i]
-    end
-    @test b == X[:,end]
-end
-
-@testset "eachbatch" begin
-    @test collect(eachbatch(X, 2)) == collect(eachbatch(X, size=2))
-
-    for (i, x) in enumerate(eachbatch(X, 2, partial=true))
-        if i != 8
-            @test size(x) == (4,2)
-            @test x == X[:,2i-1:2i]
-        else
-            @test size(x) == (4,1)
-            @test x == X[:,2i-1:2i-1]
-        end
-    end
-
-    for (i, x) in enumerate(eachbatch(X, 2, buffer=true, partial=false))
-        @test size(x) == (4,2)
-        @test x == X[:,2i-1:2i]
-    end
-
-    b = zeros(4, 2)
-    for (i, x) in enumerate(eachbatch(X, 2, buffer=b, partial=false))
-        @test size(x) == (4,2)
-        @test x == X[:,2i-1:2i]
-    end
-end
\ No newline at end of file
diff --git a/test/dataloader.jl b/test/dataloader.jl
index 1c7519e..1454079 100644
--- a/test/dataloader.jl
+++ b/test/dataloader.jl
@@ -1,84 +1,100 @@
 
 @testset "DataLoader" begin
-    X = reshape([1:10;], (2, 5))
-    Y = [1:5;]
+    X2 = reshape([1:10;], (2, 5))
+    Y2 = [1:5;]
 
-    d = DataLoader(X, batchsize=2)
-    @inferred first(d)
+    d = DataLoader(X2, batchsize=2)
+    @test_broken @inferred(first(d)) isa Array
     batches = collect(d)
-    @test eltype(batches) == eltype(d) == typeof(X)
+    @test_broken  eltype(d) == typeof(X2)
+    @test eltype(batches) == typeof(X2)
     @test length(batches) == 3
-    @test batches[1] == X[:,1:2]
-    @test batches[2] == X[:,3:4]
-    @test batches[3] == X[:,5:5]
+    @test batches[1] == X2[:,1:2]
+    @test batches[2] == X2[:,3:4]
+    @test batches[3] == X2[:,5:5]
 
-    d = DataLoader(X, batchsize=2, partial=false)
-    @inferred first(d)
+    d = DataLoader(X2, batchsize=2, partial=false)
+    # @inferred first(d)
     batches = collect(d)
-    @test eltype(batches) == eltype(d) == typeof(X)
+    @test_broken eltype(d) == typeof(X2)
     @test length(batches) == 2
-    @test batches[1] == X[:,1:2]
-    @test batches[2] == X[:,3:4]
+    @test batches[1] == X2[:,1:2]
+    @test batches[2] == X2[:,3:4]
 
-    d = DataLoader((X,), batchsize=2, partial=false)
-    @inferred first(d)
+    d = DataLoader((X2,), batchsize=2, partial=false)
+    # @inferred first(d)
     batches = collect(d)
-    @test eltype(batches) == eltype(d) == Tuple{typeof(X)}
+    @test_broken eltype(d) == Tuple{typeof(X2)}
+    @test eltype(batches) == Tuple{typeof(X2)}
     @test length(batches) == 2
-    @test batches[1] == (X[:,1:2],)
-    @test batches[2] == (X[:,3:4],)
+    @test batches[1] == (X2[:,1:2],)
+    @test batches[2] == (X2[:,3:4],)
 
-    d = DataLoader((X, Y), batchsize=2)
-    @inferred first(d)
+    d = DataLoader((X2, Y2), batchsize=2)
+    # @inferred first(d)
     batches = collect(d)
-    @test eltype(batches) == eltype(d) == Tuple{typeof(X), typeof(Y)}
+    @test_broken eltype(d) == Tuple{typeof(X2), typeof(Y2)}
+    @test eltype(batches) == Tuple{typeof(X2), typeof(Y2)}
     @test length(batches) == 3
     @test length(batches[1]) == 2
     @test length(batches[2]) == 2
     @test length(batches[3]) == 2
-    @test batches[1][1] == X[:,1:2]
-    @test batches[1][2] == Y[1:2]
-    @test batches[2][1] == X[:,3:4]
-    @test batches[2][2] == Y[3:4]
-    @test batches[3][1] == X[:,5:5]
-    @test batches[3][2] == Y[5:5]
+    @test batches[1][1] == X2[:,1:2]
+    @test batches[1][2] == Y2[1:2]
+    @test batches[2][1] == X2[:,3:4]
+    @test batches[2][2] == Y2[3:4]
+    @test batches[3][1] == X2[:,5:5]
+    @test batches[3][2] == Y2[5:5]
 
     # test with NamedTuple
-    d = DataLoader((x=X, y=Y), batchsize=2)
-    @inferred first(d)
+    d = DataLoader((x=X2, y=Y2), batchsize=2)
+    # @inferred first(d)
     batches = collect(d)
-    @test eltype(batches) == eltype(d) == NamedTuple{(:x, :y), Tuple{typeof(X), typeof(Y)}}
+    @test_broken eltype(d) == NamedTuple{(:x, :y), Tuple{typeof(X2), typeof(Y2)}}
+    @test eltype(batches) == NamedTuple{(:x, :y), Tuple{typeof(X2), typeof(Y2)}}
     @test length(batches) == 3
     @test length(batches[1]) == 2
     @test length(batches[2]) == 2
     @test length(batches[3]) == 2
-    @test batches[1][1] == batches[1].x == X[:,1:2]
-    @test batches[1][2] == batches[1].y == Y[1:2]
-    @test batches[2][1] == batches[2].x == X[:,3:4]
-    @test batches[2][2] == batches[2].y == Y[3:4]
-    @test batches[3][1] == batches[3].x == X[:,5:5]
-    @test batches[3][2] == batches[3].y == Y[5:5]
+    @test batches[1][1] == batches[1].x == X2[:,1:2]
+    @test batches[1][2] == batches[1].y == Y2[1:2]
+    @test batches[2][1] == batches[2].x == X2[:,3:4]
+    @test batches[2][2] == batches[2].y == Y2[3:4]
+    @test batches[3][1] == batches[3].x == X2[:,5:5]
+    @test batches[3][2] == batches[3].y == Y2[5:5]
 
-    # test interaction with `train!`
-    θ = ones(2)
-    X = zeros(2, 10)
-    d  = DataLoader(X)
+    # test iteration
+    X3 = zeros(2, 10)
+    d  = DataLoader(X3)
     for x in d
         @test size(x) == (2,1)
     end
     
-    # test interaction with `train!`
-    θ = zeros(2)
-    X = ones(2, 10)
-    Y = fill(5, 10)
-    d  = DataLoader((X, Y))
+    # test iteration
+    X3 = ones(2, 10)
+    Y3 = fill(5, 10)
+    d  = DataLoader((X3, Y3))
     for (x, y) in d
         @test size(x) == (2,1)
         @test y == [5]
     end
-    # specify the rng
-    d = map(identity, DataLoader(X, batchsize=2; shuffle=true, rng=Random.seed!(Random.default_rng(), 5)))
 
+    @testset "shuffle & rng" begin
+        X4 = rand(2, 1000)
+        d1 = DataLoader(X4, batchsize=2; shuffle=true)
+        d2 = DataLoader(X4, batchsize=2; shuffle=true)
+        @test first(d1) != first(d2)
+        Random.seed!(17)
+        d1 = DataLoader(X4, batchsize=2; shuffle=true)
+        x1 = first(d1)
+        Random.seed!(17)
+        d2 = DataLoader(X4, batchsize=2; shuffle=true)
+        @test x1 == first(d2)
+        d1 = DataLoader(X4, batchsize=2; shuffle=true, rng=MersenneTwister(1))
+        d2 = DataLoader(X4, batchsize=2; shuffle=true, rng=MersenneTwister(1))
+        @test first(d1) == first(d2)
+    end
+    
     # numobs/getobs compatibility
     d = DataLoader(CustomType(), batchsize=2)
     @test first(d) == [1, 2]
@@ -87,13 +103,29 @@
     @testset "Dict" begin
         data = Dict("x" => rand(2,4), "y" => rand(4))
         dloader = DataLoader(data, batchsize=2)
-        @test eltype(dloader) == Dict{String, Array{Float64}}
+        @test_broken eltype(dloader) == Dict{String, Array{Float64}}
         c = collect(dloader)
+        @test eltype(c) == Dict{String, Array{Float64}}
         @test c[1] == Dict("x" => data["x"][:,1:2], "y" => data["y"][1:2])
         @test c[2] == Dict("x" => data["x"][:,3:4], "y" => data["y"][3:4])
 
         data = Dict("x" => rand(2,4), "y" => rand(2,4))
         dloader = DataLoader(data, batchsize=2)
-        @test eltype(dloader) == Dict{String, Matrix{Float64}}
+        @test_broken eltype(dloader) == Dict{String, Matrix{Float64}}
+        @test eltype(collect(dloader)) == Dict{String, Matrix{Float64}}
+    end
+
+
+    @testset "range" begin
+        data = 1:10
+
+        dloader = DataLoader(data, batchsize=2)
+        c = collect(dloader)
+        @test eltype(c) == UnitRange{Int64}
+        @test c[1] == 1:2
+
+        dloader = DataLoader(data, batchsize=2, shuffle=true)
+        c = collect(dloader)
+        @test eltype(c) == Vector{Int}
     end
-end
\ No newline at end of file
+end
diff --git a/test/eachobs.jl b/test/eachobs.jl
new file mode 100644
index 0000000..2b4d419
--- /dev/null
+++ b/test/eachobs.jl
@@ -0,0 +1,38 @@
+@testset "eachobs" begin
+    for (i,x) in enumerate(eachobs(X))
+        @test x == X[:,i]
+    end
+
+    for (i,x) in enumerate(eachobs(X, buffer=true))
+        @test x == X[:,i]
+    end
+
+    b = zeros(size(X, 1))
+    for (i,x) in enumerate(eachobs(X, buffer=b))
+        @test x == X[:,i]
+    end
+    @test b == X[:,end]
+
+    @testset "batched" begin
+        for (i, x) in enumerate(eachobs(X, batchsize=2, partial=true))
+            if i != 8
+                @test size(x) == (4,2)
+                @test x == X[:,2i-1:2i]
+            else
+                @test size(x) == (4,1)
+                @test x == X[:,2i-1:2i-1]
+            end
+        end
+
+        for (i, x) in enumerate(eachobs(X, batchsize=2, buffer=true, partial=false))
+            @test size(x) == (4,2)
+            @test x == X[:,2i-1:2i]
+        end
+
+        b = zeros(4, 2)
+        for (i, x) in enumerate(eachobs(X, batchsize=2, buffer=b, partial=false))
+            @test size(x) == (4,2)
+            @test x == X[:,2i-1:2i]
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 24257bd..78f8762 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -38,7 +38,7 @@ MLUtils.getobs(::CustomType, i::AbstractVector) = collect(i)
 # @testset "MLUtils.jl" begin
 
 @testset "batchview" begin; include("batchview.jl"); end
-@testset "dataiterator" begin; include("dataiterator.jl"); end
+@testset "eachobs" begin; include("eachobs.jl"); end
 @testset "dataloader" begin; include("dataloader.jl"); end
 @testset "folds" begin; include("folds.jl"); end
 @testset "observation" begin; include("observation.jl"); end