Add method for chunk with size of chunks as kwarg

Author: theabhirath

src/utils.jl

@@ -158,28 +158,27 @@ julia> xs[1]
  3  8  13  18
 ```
 """
-chunk(x, n::Int) = collect(Iterators.partition(x, ceil(Int, length(x) / n)))
+chunk(x; size::Int) = collect(Iterators.partition(x, size))
+chunk(x, n::Int) = chunk(x; size = ceil(Int, length(x) / n))
 
-function chunk(x::AbstractArray, n::Int; dims::Int=ndims(x))
-    idxs = _partition_idxs(x, n, dims) 
+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)
 
-function _partition_idxs(x, n, dims)
-    bs = ceil(Int, size(x, dims) / n)
-    Iterators.partition(axes(x, dims), bs)
-end
-
-function rrule(::typeof(chunk), x::AbstractArray, n::Int; dims::Int=ndims(x))
+function rrule(::typeof(chunk), x::AbstractArray; size::Int, dims::Int=ndims(x))
     # this is the implementation of chunk
-    idxs = _partition_idxs(x, n, dims) 
+    idxs = _partition_idxs(x, size, dims)
     y = [selectdim(x, dims, i) for i in idxs]
     valdims = Val(dims)
-    chunk_pullback(dy) = (NoTangent(), ∇chunk(unthunk(dy), x, idxs, valdims), NoTangent())
-    
+    chunk_pullback(dy) = (NoTangent(), ∇chunk(unthunk(dy), x, idxs, valdims))
+
     return y, chunk_pullback
 end
 
+_partition_idxs(x, size, dims) = Iterators.partition(axes(x, dims), size)
+
 # Similar to ∇eachslice  https://github.com/JuliaDiff/ChainRules.jl/blob/8108a77a96af5d4b0c460aac393e44f8943f3c5e/src/rulesets/Base/indexing.jl#L77
 function ∇chunk(dys, x::AbstractArray, idxs, vd::Val{dim}) where {dim}
     i1 = findfirst(dy -> !(dy isa AbstractZero), dys)

test/utils.jl

@@ -111,10 +111,10 @@ end
     n = 2
     dims = 2
     x = rand(4, 5)
-    y = chunk(x, 2)
-    dy = randn!.(collect.(y))
-    idxs = MLUtils._partition_idxs(x, n, dims)
-    test_zygote(MLUtils.∇chunk, dy, x, idxs, Val(dims), check_inferred=false)
+    l = chunk(x, 2)
+    dl = randn!.(collect.(l))
+    idxs = MLUtils._partition_idxs(x, ceil(Int, size(x, dims) / n), dims)
+    test_zygote(MLUtils.∇chunk, dl, x, idxs, Val(dims), check_inferred=false)
 end
 
 @testset "group_counts" begin