Use cld and fld

Add tests for `chunk` with `size`

Author: theabhirath

src/batchview.jl

@@ -100,7 +100,7 @@ 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)
-    count = partial ? ceil(Int, n / batchsize) : floor(Int, n / batchsize)
+    count = partial ? cld(n, batchsize) : fld(n, batchsize)
     BatchView{E,T,typeof(collate)}(data, batchsize, count, partial)
 end
 

src/utils.jl

@@ -181,13 +181,13 @@ julia> xes[2]
 ```
 """
 chunk(x; size::Int) = collect(Iterators.partition(x, size))
-chunk(x, n::Int) = chunk(x; size = ceil(Int, length(x) / n))
+chunk(x, n::Int) = chunk(x; size = cld(length(x), n))
 
 function chunk(x::AbstractArray; size::Int, dims::Int=ndims(x))
     idxs = _partition_idxs(x, size, dims)
     [selectdim(x, dims, i) for i in idxs]
 end
-chunk(x::AbstractArray, n::Int; dims::Int=ndims(x)) = chunk(x; size = ceil(Int, size(x, dims) / n), dims)
+chunk(x::AbstractArray, n::Int; dims::Int=ndims(x)) = chunk(x; size = cld(size(x, dims), n), dims)
 
 function rrule(::typeof(chunk), x::AbstractArray; size::Int, dims::Int=ndims(x))
     # this is the implementation of chunk

test/utils.jl

@@ -101,9 +101,16 @@ end
     x = reshape(collect(1:20), (5, 4))
     cs = chunk(x, 2)
     @test length(cs) == 2
-    cs[1] == [1  6; 2  7; 3  8; 4  9; 5 10]
-    cs[2] == [11 16; 12 17; 13 18; 14 19; 15 20]
-    
+    @test cs[1] == [1  6; 2  7; 3  8; 4  9; 5 10]
+    @test cs[2] == [11 16; 12 17; 13 18; 14 19; 15 20]
+
+    x = permutedims(reshape(collect(1:10), (2, 5)))
+    cs = chunk(x; size = 2, dims = 1)
+    @test length(cs) == 3
+    @test cs[1] == [1 2; 3 4]
+    @test cs[2] == [5 6; 7 8]
+    @test cs[3] == [9 10]
+
     # test gradient
     test_zygote(chunk, rand(10), 3, check_inferred=false)
 
@@ -113,7 +120,7 @@ end
     x = rand(4, 5)
     l = chunk(x, 2)
     dl = randn!.(collect.(l))
-    idxs = MLUtils._partition_idxs(x, ceil(Int, size(x, dims) / n), dims)
+    idxs = MLUtils._partition_idxs(x, cld(size(x, dims), n), dims)
     test_zygote(MLUtils.∇chunk, dl, x, idxs, Val(dims), check_inferred=false)
 end