iterator over columns of high-dimensional array

columns.jl

"""
    columns(A)

Returns an array over the columns of `A` (as views). Similar to `eachcol` but
for higher-dimensional arrays. In general column (i,j,k,...) is defined as
`A[:,i,j,k,...]`.
"""
function columns(A::AbstractArray)
    [A[:,I] for I in CartesianIndices(Base.tail(axes(A)))]
end

@testset "columns" begin
    A = randn(4,5)
    @test columns(A) == collect(eachcol(A))
    @inferred columns(A)
    A = randn(5,5,3)
    @test vec(columns(A)) == collect(eachcol(reshape(A,5,15)))
    @inferred columns(A)
end