Extend chunk to take a Vector of sizes

[theabhirath]

The last feature we need to achieve parity with PyTorch would be to provide a dispatch of the type chunk(x; sizes::Vector{Int}), which would take in a Vector of integers and split the input into length(sizes) chunks, with the sizes of each chunk according to sizes. This should achieve something like:

julia> a = reshape(collect(1:10), (5, 2))
5×2 Matrix{Int64}:
 1   6
 2   7
 3   8
 4   9
 5  10
julia> collect(chunk(a; dims = 1, sizes = [1, 4]))
[1  6]
[2 7; 3 8; 4 9; 5 10]

[ToucheSir]

Is there a benefit to making this vector over tuple?

[theabhirath]

I don't think so, no. Either would be fine. Is there a generic container type that handles both or would have to write out a Union?

[ToucheSir]

Probably a union. I think using generic functions like map would let you share an implementation as well, but I have not tried.