From e41826cc1fc7dc4181e507a149dab0635d3624de Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:23 +0200 Subject: [PATCH 01/45] Require Julia v1.10, update CI --- .github/workflows/{ci.yml => CI.yml} | 2 +- .gitignore | 2 ++ Project.toml | 5 +---- src/ArraysOfArrays.jl | 6 ------ 4 files changed, 4 insertions(+), 11 deletions(-) rename .github/workflows/{ci.yml => CI.yml} (99%) diff --git a/.github/workflows/ci.yml b/.github/workflows/CI.yml similarity index 99% rename from .github/workflows/ci.yml rename to .github/workflows/CI.yml index 084929f..c3b8389 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/CI.yml @@ -26,7 +26,7 @@ jobs: fail-fast: false matrix: version: - - '1.6' + - '1.10' - '1' - 'pre' os: diff --git a/.gitignore b/.gitignore index d0b53d9..a4e2d19 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ .ipynb_checkpoints .vscode Manifest.toml +Manifest-v*.toml +LocalPreferences.toml diff --git a/Project.toml b/Project.toml index 848cd29..cd3e82c 100644 --- a/Project.toml +++ b/Project.toml @@ -3,9 +3,6 @@ uuid = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" version = "0.6.6" [deps] -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] @@ -23,4 +20,4 @@ Adapt = "1, 2, 3, 4" ChainRulesCore = "1" StaticArraysCore = "1" Statistics = "1" -julia = "1.6" +julia = "1.10" diff --git a/src/ArraysOfArrays.jl b/src/ArraysOfArrays.jl index 1828f35..8f6ccfe 100644 --- a/src/ArraysOfArrays.jl +++ b/src/ArraysOfArrays.jl @@ -21,10 +21,4 @@ include("array_of_similar_arrays.jl") include("vector_of_arrays.jl") include("broadcasting.jl") -@static if !isdefined(Base, :get_extension) - include("../ext/ArraysOfArraysAdaptExt.jl") - include("../ext/ArraysOfArraysChainRulesCoreExt.jl") - include("../ext/ArraysOfArraysStaticArraysCoreExt.jl") -end - end # module From 5cddfa1561c7de2a80c52f18cfd6fc2cd6706749 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 02/45] Improve implementation of _convert_elype --- src/util.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/util.jl b/src/util.jl index 10ae248..e15a0ce 100644 --- a/src/util.jl +++ b/src/util.jl @@ -32,7 +32,7 @@ Base.@propagate_inbounds swap_front_back_tuple(x::NTuple{N,Any}, ::Val{M}) where _convert_elype(::Type{T}, A::AbstractArray{T}) where {T} = A -_convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(x -> convert(T, x), A) +_convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1(convert, T), A) Base.@pure _add_vals(Val_M::Val{M}, Val_N::Val{N}) where {M,N} = From a55442aeb2ef7fd16a0d8d205b27876487639752 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 03/45] Improve implementation of _add_vals --- src/util.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/util.jl b/src/util.jl index e15a0ce..8baccf8 100644 --- a/src/util.jl +++ b/src/util.jl @@ -35,9 +35,7 @@ _convert_elype(::Type{T}, A::AbstractArray{T}) where {T} = A _convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1(convert, T), A) -Base.@pure _add_vals(Val_M::Val{M}, Val_N::Val{N}) where {M,N} = - Val{length((ntuple(identity, Val_M)..., ntuple(identity, Val_N)...))}() - +Base.@pure _add_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A + B}() Base.@pure require_ndims(A::AbstractArray{T,N}, Val_N::Val{N}) where {T,N} = nothing From 4fd3b254e8cfa4cf4e5e46054621f8a51437d66e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 04/45] Disable flatview for generic arrays of arrays Breaking, but current fallback to Iterators.flatten caused more trouble than it was worth. --- src/functions.jl | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/functions.jl b/src/functions.jl index 7509e9d..606eef8 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -36,8 +36,7 @@ deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = View array `A` in a suitable flattened form. The shape of the flattened form will depend on the type of `A`. If the `A` is not a nested array, the return -value is `A` itself. When no type-specific method is available, `flatview` -will fall back to `Base.Iterators.flatten`. +value is `A` itself. Only specific types of nested arrays are supported. """ function flatview end export flatview @@ -45,8 +44,8 @@ export flatview @inline flatview(A::AbstractArray) = A # TODO: Implement flatview on generic nested arrays via new `FlatView`, using -# deepgetindex to implement getindex, etc. -@inline flatview(A::AbstractArray{<:AbstractArray}) = Base.Iterators.flatten(A) +# deepgetindex to implement getindex, etc? +@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented for generic AbstractArray{<:AbstractArray}")) """ From b608d83dbabef0d24a6b76f219a1aa7bcd422e27 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 05/45] Remove deepgetindex, deepsetindex! and deepview --- docs/src/index.md | 1 - src/array_of_similar_arrays.jl | 46 -------------- src/functions.jl | 109 +------------------------------- test/array_of_similar_arrays.jl | 56 ---------------- test/functions.jl | 50 --------------- 5 files changed, 1 insertion(+), 261 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 114dfce..b4c27dc 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -6,7 +6,6 @@ This package also defines and exports the following new functions applicable to * [`nestedview`](@ref) and [`flatview`](@ref) switch between a flat and a nested view of the same data. * [`innersize`](@ref) returns the size of the elements of an array, provided they all have equal size. -* [`deepgetindex`](@ref), [`deepsetindex!`](@ref) and [`deepview`](@ref) provide index-based access across multiple layers of nested arrays * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. * [`abstract_nestedarray_type`](@ref) returns the type of nested `AbstractArray`s for a given innermost element type with multiple layers of nesting. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 7d6a0b2..ab381ec 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -217,52 +217,6 @@ Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) = Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs...))}() -Base.@propagate_inbounds function deepgetindex(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - nested = _result_is_nested(idxs_outer, idxs_inner) - _deepgetindex_impl_aosa(A, idxs_outer, idxs_inner, nested) -end - -Base.@propagate_inbounds _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - getindex(A.data, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds function _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) - new_data = getindex(A.data, idxs_inner..., idxs_outer...) - nestedview(new_data, ndims_after_getindex(idxs_inner)) -end - - -Base.@propagate_inbounds function deepsetindex!(A::ArrayOfSimilarArrays{T,M,N,L}, x, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepsetindex_impl_aosa!(A, x, idxs_outer, idxs_inner, _result_is_nested(idxs_outer, idxs_inner)) - A -end - -Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - setindex!(A.data, x, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) = - _deepsetindex_impl!(A, x, idxs_outer, idxs_inner) - -# TODO: Specialized method _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) = - - -Base.@propagate_inbounds function deepview(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - nested = _result_is_nested(idxs_outer, idxs_inner) - _deepview_impl_aosa(A, idxs_outer, idxs_inner, nested) -end - -Base.@propagate_inbounds _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - view(A.data, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds function _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) - new_data = view(A.data, idxs_inner..., idxs_outer...) - nestedview(new_data, ndims_after_getindex(idxs_inner)) -end - - - const VectorOfSimilarArrays{ T, M, L, P<:AbstractArray{T,L} diff --git a/src/functions.jl b/src/functions.jl index 606eef8..f800a8d 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -42,10 +42,7 @@ function flatview end export flatview @inline flatview(A::AbstractArray) = A - -# TODO: Implement flatview on generic nested arrays via new `FlatView`, using -# deepgetindex to implement getindex, etc? -@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented for generic AbstractArray{<:AbstractArray}")) +@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) """ @@ -110,110 +107,6 @@ end @inline innersize(ref::Ref) = size(only(ref)) -""" - deepgetindex(A::AbstractArray, idxs...) - deepgetindex(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N} - -Recursive `getindex` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepgetindex` is equivalent to -`getindex`. - -See also [`deepsetindex!`](@ref) and [`deepview`](@ref). -""" -function deepgetindex end -export deepgetindex - -Base.@propagate_inbounds deepgetindex(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = getindex(A, idxs...) -Base.@propagate_inbounds deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = getindex(A, idxs...) - -Base.@propagate_inbounds function deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepgetindex_impl(A, idxs_outer, idxs_inner) -end - -Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N} = - deepgetindex(getindex(A, idxs_outer...), idxs_inner...) - -Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N} = - _deepgetindex_tupled.(view(A, idxs_outer...), (idxs_inner,)) - -Base.@propagate_inbounds _deepgetindex_tupled(A::AbstractArray, idxs::Tuple) = deepgetindex(A, idxs...) - - -""" - deepsetindex!(A::AbstractArray, x, idxs...) - deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N} - -Recursive `setindex!` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepsetindex!` is equivalent to -`setindex!`. - -See also [`deepgetindex`](@ref) and [`deepview`](@ref). -""" -function deepsetindex! end -export deepsetindex! - -Base.@propagate_inbounds deepsetindex!(A::AbstractArray{T,N}, x, idxs::Vararg{Any,N}) where {T,N} = setindex!(A, x, idxs...) -Base.@propagate_inbounds deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs::Vararg{Any,N}) where {N} = setindex!(A, x, idxs...) - -Base.@propagate_inbounds function deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepsetindex_impl!(A, x, idxs_outer, idxs_inner) - A -end - -Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N} - B = getindex(A, idxs_outer...) - deepsetindex!(B, x, idxs_inner...) -end - -Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N} - B = view(A, idxs_outer...) - for i in eachindex(B, x) - deepsetindex!(B[i], x[i], idxs_inner...) - end -end - - - -""" - deepview(A::AbstractArray, idxs...) - deepview(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N} - -Recursive `view` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepview` is equivalent to `view`. - -See also [`deepgetindex`](@ref) and [`deepsetindex!`](@ref). -""" -function deepview end -export deepview - -Base.@propagate_inbounds deepview(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = view(A, idxs...) -Base.@propagate_inbounds deepview(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = view(A, idxs...) - -Base.@propagate_inbounds function deepview(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepview_impl(A, idxs_outer, idxs_inner) -end - -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Real}) where {N,M} = - deepview(getindex(A, idxs_outer...), idxs_inner...) - -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Any}) where {N,M} = - deepview(getindex(A, idxs_outer...), idxs_inner...) - -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Real}) where {N,M} = - throw(ArgumentError("deepview not supported yes with outer indices $idxs_outer and inner indices $idxs_inner")) - -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Any}) where {N,M} = - _deepview_tupled.(view(A, idxs_outer...), (idxs_inner,)) - -Base.@propagate_inbounds _deepview_tupled(A::AbstractArray, idxs::Tuple) = deepview(A, idxs...) - - """ abstract_nestedarray_type(T_inner::Type, ::Val{ndims_tuple}) diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 4c2ca84..7fcdffa 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -187,62 +187,6 @@ using StatsBase: cov2cor end - @testset "deepgetindex" begin - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - - @test deepgetindex(A, 3, 4, 2, 1, 2) == @inferred deepgetindex(B, 3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 2:4, 2, 1, 2) - @test deepgetindex(A, 2, 4, :, 1, 2) == @inferred deepgetindex(B, 2, 4, :, 1, 2) - @test deepgetindex(A, 2, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2, 4, :, 1, 1:2) - @test deepgetindex(A, 2:3, 4, :, 1, 2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 2) - @test deepgetindex(A, 2:3, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 1:2) - end - - - @testset "deepsetindex!" begin - function testdata() - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - A, B - end - - A, B = testdata() - @test deepsetindex!(A, 42, 3, 4, 2, 1, 2) == @inferred deepsetindex!(B, 42, 3, 4, 2, 1, 2) - @test deepgetindex(B, 3, 4, 2, 1, 2) == 42 - - A, B = testdata() - X1 = rand(2) - @test deepsetindex!(A, X1, 2:3, 4, 2, 1, 2) == @inferred deepsetindex!(B, X1, 2:3, 4, 2, 1, 2) - @test deepgetindex(B, 2:3, 4, 2, 1, 2) == X1 - - A, B = testdata() - X2 = rand(2,2) - @test deepsetindex!(A, X2, 2, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X2, 2, 4, :, 1, 1:2) - @test deepgetindex(B, 2, 4, :, 1, 1:2) == X2 - - A, B = testdata() - X3 = [rand(2,2), rand(2,2)] - @test deepsetindex!(A, X3, 2:3, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X3, 2:3, 4, :, 1, 1:2) - @test deepgetindex(B, 2:3, 4, :, 1, 1:2) == X3 - end - - - @testset "deepview" begin - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - - @test deepview(A, 3, 4, 2, 1, 2) == @inferred deepview(B, 3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepview(B, 2:3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepview(B, 2:3, 2:4, 2, 1, 2) - @test deepview(A, 2, 4, :, 1, 2) == @inferred deepview(B, 2, 4, :, 1, 2) - @test deepview(A, 2, 4, :, 1, 1:2) == @inferred deepview(B, 2, 4, :, 1, 1:2) - @test deepview(A, 2:3, 4, :, 1, 2) == @inferred deepview(B, 2:3, 4, :, 1, 2) - @test deepview(A, 2:3, 4, :, 1, 1:2) == @inferred deepview(B, 2:3, 4, :, 1, 1:2) - end - - @testset "empty" begin A = [rand(2,3), rand(2,3), rand(2,3)] B = ArrayOfSimilarArrays(A) diff --git a/test/functions.jl b/test/functions.jl index 6b9f2d9..2f600f5 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -35,56 +35,6 @@ using StaticArrays end - @testset "deepgetindex" begin - A = gen_nested() - @test @inferred(deepgetindex(A, 1, 2)) === A[1, 2] - @test @inferred(deepgetindex(A, 1, 2, 2, 1)) == A[1, 2][2, 1] - @test @inferred(deepgetindex(A, 1, 2, 1:2, 1)) == A[1, 2][1:2, 1] - @test @inferred(deepgetindex(A, 1, 1:2, 2, 1)) == [A[1, 1][2, 1], A[1, 2][2, 1]] - @test @inferred(deepgetindex(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]] - @test_throws MethodError deepgetindex(A, 1, 2, 2, 1, 2) - - B = rand(3, 4, 5) - @test @inferred(deepgetindex(B, 2, 3, 4)) === getindex(B, 2, 3, 4) - @test_throws MethodError deepgetindex(B, 6) - @test_throws MethodError deepgetindex(B, 2, 3) - end - - - @testset "deepsetindex!" begin - B12 = [20 21 22; 23 24 25] - X = [[41, 42], [43, 44]] - - A = gen_nested() - @test @inferred(deepsetindex!(A, B12, 1, 2)) === A - @test A[1, 2] == B12 - - A = gen_nested() - @test @inferred(deepsetindex!(A, 42, 1, 2, 2, 1)) === A - @test A[1, 2][2, 1] == 42 - - A = gen_nested() - @test @inferred(deepsetindex!(A, X, 1, 1:2, 1:2, 2)) === A - @test deepgetindex(A, 1, 1:2, 1:2, 2) == X - end - - - @testset "deepview" begin - A = gen_nested() - @test @inferred(deepview(A, 1, 2)) == view(A, 1, 2) - @test @inferred(deepview(A, 1, 2, 2, 1)) == view(A[1, 2], 2, 1) - @test @inferred(deepview(A, 1, 2, 1:2, 1)) == view(A[1, 2], 1:2, 1) - @test_throws ArgumentError deepview(A, 1, 1:2, 2, 1) - @test @inferred(deepview(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]] - @test_throws MethodError deepview(A, 1, 2, 2, 1, 2) - - B = rand(3, 4, 5) - @test @inferred(deepview(B, 2, 3, 4)) === view(B, 2, 3, 4) - @test_throws MethodError deepview(B, 6) - @test_throws MethodError deepview(B, 2, 3) - end - - @testset "abstract_nestedarray_type" begin @test @inferred(abstract_nestedarray_type(Int, Val(()))) == Int @test @inferred(abstract_nestedarray_type(Int, Val((2,)))) == AbstractArray{Int, 2} From 79bdc73fec74962eb56d502065f12a0bc904f97a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 06/45] Support AbstractSlices, add getslicemap --- src/functions.jl | 24 ++++++++++++++++++++++++ src/util.jl | 36 ++++++++++++++++++++++++++++++++++++ test/functions.jl | 15 +++++++++++++++ 3 files changed, 75 insertions(+) diff --git a/src/functions.jl b/src/functions.jl index f800a8d..29d7648 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -30,6 +30,20 @@ deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = map(X -> deepmap(f, X), A) +""" + getslicemap(A::AbstractSlices) + +Get the slicemap of an `A`. + +A slicemap must have type `Tuple{Vararg{Union{Colon,Integer}}}`, e.g. +`(:, :, :, 1, 2)` or `(:, 2, :, 1, :)`. +""" +function getslicemap end +export getslicemap + +getslicemap(A::Slices) = A.slicemap + + """ flatview(A::AbstractArray) flatview(A::AbstractArray{<:AbstractArray{<:...}}) @@ -44,6 +58,14 @@ export flatview @inline flatview(A::AbstractArray) = A @inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) +function flatview(A::AbstractSlices) + if _is_aoa_slicemap(A.slicemap) + return parent(A) + else + throw(ArgumentError("flatview for AbstractSlices requires inner dimensions to be first and no dimension reordering, but slicemap is $(A.slicemap)")) + end +end + """ nestedview(A::AbstractArray{T,M+N}, M::Integer) @@ -100,6 +122,8 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} s end +@inline innersize(A::AbstractSlices) = _extract_innerdims(size(parent(A)), getslicemap(A)) + @inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = innersize(A)[dim] diff --git a/src/util.jl b/src/util.jl index 8baccf8..f0f02bd 100644 --- a/src/util.jl +++ b/src/util.jl @@ -15,6 +15,10 @@ end Base.@pure _ncolons(::Val{N}) where N = ntuple(_ -> Colon(), Val{N}()) +Base.@pure _nColons(::Val{N}) where N = ntuple(_ -> Colon, Val{N}()) + +@inline _oneto_tpl(::Val{N}) where N = ntuple(identity, Val{N}()) +Base.@pure _nInts(::Val{N}) where N = ntuple(_ -> Int, Val{N}()) Base.@propagate_inbounds front_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} = @@ -42,3 +46,35 @@ Base.@pure require_ndims(A::AbstractArray{T,N}, Val_N::Val{N}) where {T,N} = Base.@pure require_ndims(A::AbstractArray{T,M}, Val_N::Val{N}) where {T,M,N} = throw(ArgumentError("Require an array with $N dimensions")) + +@inline @generated function _extract_innerdims(obj::Tuple, slicemap::Tuple{Vararg{Union{Colon,Integer}}}) + # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), + # extract only the elements of obj where the slicemap is a Colon. + expr = Expr(:tuple) + slicepars = slicemap.parameters + for i in 1:length(slicepars) + if slicepars[i] <: Colon + push!(expr.args, :(obj[$i])) + end + end + return expr +end + +@inline @generated function _extract_outerdims(obj::Tuple, slicemap::Tuple{Vararg{Union{Colon,Integer}}}) + # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), + # extract only the elements of obj where the slicemap is a Colon. + expr = Expr(:tuple) + slicepars = slicemap.parameters + for i in 1:length(slicepars) + if slicepars[i] <: Integer + push!(expr.args, :(obj[$i])) + end + end + return expr +end + + +function _is_aoa_slicemap(slicemap::Tuple{Vararg{Union{Colon,Integer}}}) + dims = _oneto_tpl(Val(length(slicemap))) + issorted((_extract_innerdims(dims, slicemap)..., _extract_outerdims(dims, slicemap)...)) +end diff --git a/test/functions.jl b/test/functions.jl index 2f600f5..15d32ce 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -7,6 +7,9 @@ using StaticArrays @testset "functions" begin + Aes1 = eachslice(rand(5,6,7,8,9); dims = (4,5)) + Aes2 = eachslice(rand(5,6,7,8,9); dims = (4,2)) + function gen_nested() A11 = [1 2; 3 4] A21 = [4 5 6; 7 8 9] @@ -22,6 +25,15 @@ using StaticArrays B = rand(3, 2, 4) @test @inferred(nestedview(flatview(B), SVector{3})) == @inferred(nestedview(B, Val(1))) + + @test @inferred flatview(Aes1) == parent(Aes1) + @test_throws ArgumentError @inferred flatview(Aes2) + end + + + @testset "getslicemap" begin + @test @inferred(getslicemap(Aes1)) == (:, :, :, 1, 2) + @test @inferred(getslicemap(Aes2)) == (:, 2, :, 1, :) end @@ -32,6 +44,9 @@ using StaticArrays @test @inferred(innersize((2:5,))) == (4,) @test @inferred(innersize(Ref(2:5))) == (4,) @test_throws DimensionMismatch @inferred(innersize([[1, 2, 3], [4, 5]])) + + @test @inferred(innersize(Aes1)) == (5, 6, 7) + @test @inferred(innersize(Aes2)) == (5, 7, 9) end From 47546c7a8ba855d4bd33d2d2e19a41d1d65571f9 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 07/45] Make AbstractArrayOfSimilarArrays a subtype of AbstractSlices --- ext/ArraysOfArraysChainRulesCoreExt.jl | 2 +- src/array_of_similar_arrays.jl | 89 +++++++++++++++----------- test/array_of_similar_arrays.jl | 76 ++++++++++++++-------- 3 files changed, 104 insertions(+), 63 deletions(-) diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index 6d2e143..e4f9055 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -13,7 +13,7 @@ function _aosa_ctor_fromflat_pullback(ΔΩ) NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) end -function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U} +function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback end diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index ab381ec..b3a5138 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -1,7 +1,7 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). """ - AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{<:AbstractArray{T,M},N} + AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T}} <: AbstractSlices{ET,N} An array that contains arrays that have the same size/axes. The array is internally stored in flattened form as some kind of array of dimension @@ -10,31 +10,32 @@ internally stored in flattened form as some kind of array of dimension Subtypes must implement (in addition to typical array operations): flatview(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N} + getslicemap(A::SomeArrayOfSimilarArrays) The following type aliases are defined: -* `AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `AbstractVectorOfSimilarArrays{T,M,ET} = AbstractArrayOfSimilarArrays{T,M,1,ET}` +* `AbstractArrayOfSimilarVectors{T,N,ET} = AbstractArrayOfSimilarArrays{T,1,N,ET}` +* `AbstractVectorOfSimilarVectors{T,ET} = AbstractArrayOfSimilarArrays{T,1,1,ET}` """ -abstract type AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{Array{T,M},N} end +abstract type AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T}} <: AbstractSlices{ET,N} end export AbstractArrayOfSimilarArrays -const AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1} +const AbstractVectorOfSimilarArrays{T,M,ET<:AbstractArray{T}} = AbstractArrayOfSimilarArrays{T,M,1,ET} export AbstractVectorOfSimilarArrays -const AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N} +const AbstractArrayOfSimilarVectors{T,N,ET<:AbstractVector{T}} = AbstractArrayOfSimilarArrays{T,1,N,ET} export AbstractArrayOfSimilarVectors -const AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1} +const AbstractVectorOfSimilarVectors{T,ET<:AbstractVector{T}} = AbstractArrayOfSimilarArrays{T,1,1,ET} export AbstractVectorOfSimilarVectors """ - ArrayOfSimilarArrays{T,M,N,L,P} <: AbstractArrayOfSimilarArrays{T,M,N} + ArrayOfSimilarArrays{T,M,N,P} <: AbstractArrayOfSimilarArrays{T,M,N} -Represents a view of an array of dimension `L = M + N` as an array of +Represents a view of an array of dimension `M + N` as an array of dimension M with elements that are arrays with dimension N. All element arrays implicitly have equal size/axes. @@ -64,20 +65,22 @@ flatview(A_nested) === A_flat ``` """ struct ArrayOfSimilarArrays{ - T, M, N, L, - P<:AbstractArray{T,L} -} <: AbstractArrayOfSimilarArrays{T,M,N} + T, M, N, + P<:AbstractArray{T}, + ET<:AbstractArray{T,M} +} <: AbstractArrayOfSimilarArrays{T,M,N,ET} data::P - function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U,L}) where {T,M,N,L,U} + function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U}) where {T,M,N,U} require_ndims(flat_data, _add_vals(Val{M}(), Val{N}())) conv_parent = _convert_elype(T, flat_data) P = typeof(conv_parent) - new{T,M,N,L,P}(conv_parent) + ET = Base.promote_op(view, P, _nColons(Val{M}())..., _nInts(Val{N}())...) + new{T,M,N,P,ET}(conv_parent) end end -function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U,L}) where {T,M,L,U} +function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U}) where {T,M,U} _, size_outer = split_tuple(size(flat_data), Val{M}()) N = length(size_outer) ArrayOfSimilarArrays{T,M,N}(flat_data) @@ -97,13 +100,24 @@ ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(A) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M}}, flat_data::AbstractArray{U,L}) where {T,M,L,U} = R(flat_data) +Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} = R(flat_data) +Base.convert(R::Type{ArrayOfSimilarArrays{T,M}}, flat_data::AbstractArray{U}) where {T,M,U} = R(flat_data) Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) +Base.parent(A::ArrayOfSimilarArrays) = A.data +Base.stack(A::ArrayOfSimilarArrays) = A.data + +function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} + new_ET = Base.promote_op(similar, ET) + return Array{new_ET,N}(A) +end + +function getslicemap(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} + return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) +end @inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} front_tuple(size(A.data), Val{M}()) @@ -121,9 +135,9 @@ import Base.== """ - flatview(A::ArrayOfSimilarArrays{T,M,N,L,P})::P + flatview(A::ArrayOfSimilarArrays{T,M,N,P})::P -Returns the array of dimensionality `L = M + N` wrapped by `A`. The shape of +Returns the array of dimensionality `M + N` wrapped by `A`. The shape of the result may be freely changed without breaking the inner consistency of `A`. """ @@ -218,16 +232,17 @@ Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs const VectorOfSimilarArrays{ - T, M, L, - P<:AbstractArray{T,L} -} = ArrayOfSimilarArrays{T,M,1,L,P} + T, M, + P<:AbstractArray{T}, + ET<:AbstractArray{T} +} = ArrayOfSimilarArrays{T,M,1,P,ET} export VectorOfSimilarArrays -VectorOfSimilarArrays{T}(flat_data::AbstractArray{U,L}) where {T,U,L} = +VectorOfSimilarArrays{T}(flat_data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) -VectorOfSimilarArrays(flat_data::AbstractArray{T,L}) where {T,L} = +VectorOfSimilarArrays(flat_data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = @@ -237,8 +252,8 @@ VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = VectorOfSimilarArrays{T,M}(A) -Base.convert(R::Type{VectorOfSimilarArrays{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarArrays}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data) +Base.convert(R::Type{VectorOfSimilarArrays{T}}, flat_data::AbstractArray{U}) where {T,U} = R(flat_data) +Base.convert(R::Type{VectorOfSimilarArrays}, flat_data::AbstractArray{T}) where {T} = R(flat_data) Base.convert(R::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = R(A) Base.convert(R::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = R(A) @@ -289,16 +304,17 @@ end const ArrayOfSimilarVectors{ - T, N, L, - P<:AbstractArray{T,L} -} = ArrayOfSimilarArrays{T,1,N,L,P} + T, N, + P<:AbstractArray{T}, + ET<:AbstractVector{T} +} = ArrayOfSimilarArrays{T,1,N,P,ET} export ArrayOfSimilarVectors -ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U,L}) where {T,U,L} = +ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) -ArrayOfSimilarVectors(flat_data::AbstractArray{T,L}) where {T,L} = +ArrayOfSimilarVectors(flat_data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = @@ -308,16 +324,17 @@ ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = ArrayOfSimilarVectors{T,N}(A) -Base.convert(R::Type{ArrayOfSimilarVectors{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarVectors}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data) +Base.convert(R::Type{ArrayOfSimilarVectors{T}}, flat_data::AbstractArray{U}) where {T,U} = R(flat_data) +Base.convert(R::Type{ArrayOfSimilarVectors}, flat_data::AbstractArray{T}) where {T} = R(flat_data) Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(A) Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(A) const VectorOfSimilarVectors{ T, - P<:AbstractArray{T,2} -} = ArrayOfSimilarArrays{T,1,1,2,P} + P<:AbstractArray{T,2}, + ET<:AbstractVector{T} +} = ArrayOfSimilarArrays{T,1,1,P,ET} export VectorOfSimilarVectors diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 7fcdffa..75a38ea 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -38,8 +38,19 @@ using StatsBase: cov2cor } @testset "$TT from Array{Float64,$L}" begin A = rand_flat_array(Val_L) - @test typeof(@inferred TT(A)) == RT - @test typeof(@inferred convert(TT, A)) == RT + @test typeof(@inferred TT(A)) <: RT + @test typeof(@inferred convert(TT, A)) <: RT + + AosA = TT(A) + @test typeof(AosA) == typeof(convert(TT, A)) + @test @inferred(eltype(AosA)) == typeof(AosA[1]) + @test @inferred(parent(AosA)) === @inferred(flatview(AosA)) + if eltype(eltype(AosA)) == eltype(A) + @test @inferred(flatview(AosA)) === A + else + @test @inferred(flatview(AosA)) ≈ A + end + @test @inferred(stack(AosA)) === flatview(AosA) end end @@ -57,11 +68,11 @@ using StatsBase: cov2cor A_U = Array{Array{U,M},N}(A) - @test typeof(A2_ctor) == RT + @test typeof(A2_ctor) <: RT @test A2_ctor == A_U A2_conv = @inferred convert(TT, A) - @test typeof(A2_conv) == RT + @test typeof(A2_conv) <: RT @test A2_conv == A2_ctor U = eltype(flatview(A2_ctor)) @@ -73,17 +84,17 @@ using StatsBase: cov2cor @testset "construct/convert from flat array" begin - test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(5)) - test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,3,Array{Float64,3}}, Val(3)) - test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,3,Array{Float64,3}}, Val(3)) + test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,Array{Float64,3}}, Val(3)) + test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) - test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(5)) - test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,3,Array{Float32,3}}, Val(3)) - test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,3,Array{Float32,3}}, Val(3)) + test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,Array{Float32,3}}, Val(3)) + test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) @@ -92,24 +103,24 @@ using StatsBase: cov2cor @testset "construct/convert from nested arrays" begin - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,4,5,Array{Float32,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,4,Array{Float32,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(ArrayOfSimilarVectors{Float64,4}, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,4,5,Array{Float32,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float64,4}, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,4,Array{Float32,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) test_from_nested(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) test_from_nested(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(1), Val(1)) @@ -121,6 +132,19 @@ using StatsBase: cov2cor test_rrule(ArrayOfSimilarArrays{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) end + @testset "AbstractSlices interface" begin + @test @inferred(ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices + let A = ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7)) + @test @inferred(A[2, 3, 4]) isa AbstractArray{Float32,2} + ref_ET = typeof(A[2, 3, 4]) + @test A isa AbstractSlices{ref_ET, 3} + @test @inferred(eltype(A)) == ref_ET + @test @inferred(innersize(A)) == (3, 4) + @test @inferred(getslicemap(A)) == (:, :, 1, 2, 3) + @test @inferred(parent(A)) === A.data + end + end + @testset "add remove" begin EA_ref1 = rand_flat_array(Val(3)) EA_ref2 = rand_flat_array(Val(3)) From c4df3feaf14167f99be219c16c999e13e63c1efd Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 08/45] STASH partview API --- ext/ArraysOfArraysChainRulesCoreExt.jl | 42 ++++++++- src/ArraysOfArrays.jl | 1 + src/array_of_similar_arrays.jl | 1 + src/base_slices.jl | 73 ++++++++++++++++ src/functions.jl | 113 +++++++++++++++++-------- src/util.jl | 31 ------- 6 files changed, 195 insertions(+), 66 deletions(-) create mode 100644 src/base_slices.jl diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index e4f9055..a643ad3 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -2,13 +2,51 @@ module ArraysOfArraysChainRulesCoreExt -import ChainRulesCore -using ChainRulesCore: NoTangent, unthunk +using ChainRulesCore: ChainRulesCore, NoTangent, Thunk, unthunk, @thunk, @non_differentiable using ArraysOfArrays: ArrayOfSimilarArrays using ArraysOfArrays: flatview +struct _MappedMaybeThunk{F, T} <: AbstractThunk + f::F + x::T +end +ChainRulesCore.unthunk(mt::_MappedMaybeThunk) = mt.f(unthunk(mt.x)) +mapthunk(f::F, x::T) where {F,T} = _MappedMaybeThunk{F,T}(f, x) +mapthunk(::Type{F}, x::T) where {F,T} = _MappedMaybeThunk{Type{F},T}(F, x) + + +@non_differentiable getpartmode(::Any) +@non_differentiable innersize(::Any) + + + +function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, ::Unpartitioned) + return flatview(A), _partview_pullback +end +_unpart_partview_pullback(ΔΩ) = NoTangent(), ΔΩ, NoTangent() + +function ChainRulesCore.rrule(::typeof(flatview), A::AbstractArray) + return flatview(A), _unpart_flatview_pullback +end +_unpart_flatview_pullback(ΔΩ) = NoTangent(), ΔΩ + + + +function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, partmode::AbstractSlicingMode) + return partview(A, partmode), _partview_pullback +end +_partview_pullback(ΔΩ) = NoTangent(), mapthunk(flatview, ΔΩ), NoTangent() + +function ChainRulesCore.rrule(::typeof(flatview), A::AbstractArray{<:AbstractArray}) + pmode = getpartmode(A) + return flatview(A), Base.Fix2(_flatview_pullback, pmode) +end +_flatview_pullback(ΔΩ, pmode) = NoTangent(), mapthunk(Base.Fix2(partview, pmode), ΔΩ) + + + function _aosa_ctor_fromflat_pullback(ΔΩ) NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) end diff --git a/src/ArraysOfArrays.jl b/src/ArraysOfArrays.jl index 8f6ccfe..2a0788d 100644 --- a/src/ArraysOfArrays.jl +++ b/src/ArraysOfArrays.jl @@ -17,6 +17,7 @@ using Statistics include("util.jl") include("functions.jl") +include("base_slices.jl") include("array_of_similar_arrays.jl") include("vector_of_arrays.jl") include("broadcasting.jl") diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index b3a5138..14515a4 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -119,6 +119,7 @@ function getslicemap(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) end +#!!!!! @inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} front_tuple(size(A.data), Val{M}()) end diff --git a/src/base_slices.jl b/src/base_slices.jl new file mode 100644 index 0000000..f1e2e3e --- /dev/null +++ b/src/base_slices.jl @@ -0,0 +1,73 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +""" + struct BaseSlicing{N,TPL<:Tuple{Vararg{Int,N}}} <: AbstractSlicingMode + +The partitioning mode of `Slices`. + +Constructor: + +``` +BaseSlicing(slicemap::Tuple{Vararg{Union{Colon,Int}}}}) +``` + +`slicemap` equals the `slicemap` property of `Base.Slice` objects. + +See also [`AbstractSlicingMode`](@ref). +""" +struct BaseSlicing{N,TPL<:Tuple{Vararg{Union{Colon,Int},N}}} <: AbstractSlicingMode + slicemap::TPL +end + + +@inline getpartmode(A::Slices) = BaseSlicing(A.slicemap) +@inline flatview(A::Slices) = parent(A) + +@inline stacked(A::Slices) = reshape(parent(A), (length(A), prod(size(parent(A))) ÷ length(A))) + +@inline innersize(A::AbstractSlices) = getinnnerdims(size(flatview(A)), getpartmode(A)) + + +@inline @generated function getinnnerdims(obj::Tuple, pmode::BaseSlicing{N,SliceMapT}) where {N,SliceMapT} + # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), + # extract only the elements of obj where the slicemap is a Colon. + expr = Expr(:tuple) + slicepars = SliceMapT.parameters + for i in 1:length(slicepars) + if slicepars[i] <: Colon + push!(expr.args, :(obj[$i])) + end + end + return expr +end + + +@inline @generated function getouterdims(obj::Tuple, pmode::BaseSlicing{N,SliceMapT}) where {N,SliceMapT} + # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), + # extract only the elements of obj where the slicemap is a Colon. + slicepars = SliceMapT.parameters + outdimidxs = Int[] + for i in 1:length(slicepars) + if slicepars[i] <: Integer + push!(outdimidxs, i) + end + end + + outdimidxs_expr = Expr(:tuple, outdimidxs...) + idxorder_expr = Expr(:tuple, [:(slicemap[$i]) for i in outdimidxs]...) + result_expr = Expr(:tuple, [:(obj[outdimidxs[idxorder[$i]]]) for i in eachindex(outdimidxs)]...) + + quote + slicemap = pmode.slicemap + outdimidxs = $outdimidxs_expr + idxorder = $idxorder_expr + return $result_expr + end +end + + +function is_natural_dimorder_partmode(pmode::BaseSlicing) + slicemap = pmode.slicemap + dims = _oneto_tpl(Val(length(slicemap))) + issorted((_extract_innerdims(dims, slicemap)..., _extract_outerdims(dims, slicemap)...)) +end diff --git a/src/functions.jl b/src/functions.jl index 29d7648..7389d5c 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -2,71 +2,119 @@ """ - innermap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) + abstract type AbstractPartitionMode + +Abstract supertype for array partition modes. + +Use [`getpartmode`](@ref) to get the partition mode of an array. + +See also [`partview`](@ref) and [`flatview`](@ref). +""" +abstract type AbstractPartitionMode end +export AbstractPartitionMode -Nested `map` at depth 2. Equivalent to `map(X -> map(f, X) A)`. """ -function innermap end -export innermap + struct Unpartitioned <: AbstractPartitionMode -innermap(f::Base.Callable, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = - map(X -> map(f, X), A) +The partitioning mode of unpartitioned arrays. +Constructor: `Unpartitioned()` +""" +struct Unpartitioned <: AbstractPartitionMode end +export Unpartitioned """ - deepmap(f::Base.Callable, x::Any) - deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray{<:...}}) + abstract type AbstractSlicingMode <: AbstractPartitionMode -Applies `map` at the deepest possible layer of nested arrays. If `A` is not -a nested array, `deepmap` behaves identical to `Base.map`. +Abstract supertype for array partition modes. + +Use `getpartmode` to get the partition mode of an partitioned array. """ -function deepmap end -export deepmap +abstract type AbstractSlicingMode <: AbstractPartitionMode end +export AbstractSlicingMode -deepmap(f::Base.Callable, x::Any) = map(f, x) -deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = - map(X -> deepmap(f, X), A) +""" + getpartmode(A::AbstractArray)::Unpartitioned + getpartmode(A::AbstractArray{<:AbstractArray})::AbstractSlicingMode + +Get the partitioning mode of `A`. + +`partview(flatview(A), getpartmode(A))` must equal `A`, and should have +the same type as `A` if at all possible. +`getpartmode` should be an allocation-free O(1) operation, if at all possible. """ - getslicemap(A::AbstractSlices) +function getpartmode end + +@inline getpartmode(::AbstractArray) = Unpartitioned() + +function getpartmode(A::AbstractArray{<:AbstractArray}) + throw(ArgumentError("getpartmode not implemented for nested arrays of type $(nameof(typeof(A)))")) +end -Get the slicemap of an `A`. -A slicemap must have type `Tuple{Vararg{Union{Colon,Integer}}}`, e.g. -`(:, :, :, 1, 2)` or `(:, 2, :, 1, :)`. """ -function getslicemap end -export getslicemap + partview(A::AbstractArray, pmode::AbstractSlicingMode) + +View array `A` in partitioned form, as an array of arrays. -getslicemap(A::Slices) = A.slicemap +If `A` is not a nested array return `A` itself. If `A` is a partitioned array, +return the original unpartition array. + +`partview` should be an allocation-free O(1) operation, if at all possible. + +See also [`flatview`](@ref) and [`getpartmode`](@ref). +""" +function partview end +export partview + +@inline partview(A::AbstractArray, ::Unpartitioned) = A """ flatview(A::AbstractArray) flatview(A::AbstractArray{<:AbstractArray{<:...}}) -View array `A` in a suitable flattened form. The shape of the flattened form -will depend on the type of `A`. If the `A` is not a nested array, the return -value is `A` itself. Only specific types of nested arrays are supported. +View array `A` in a suitable flattened form. + +`partview(flatview(A), getpartmode(A))` must equal `A`, and should have +the same type as `A` if at all possible. + +If `A` is not a nested array return `A` itself. If `A` is a partitioned array, +return the original unpartition array. + +`flatview` should be an allocation-free O(1) operation, if at all possible. """ function flatview end export flatview @inline flatview(A::AbstractArray) = A -@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) -function flatview(A::AbstractSlices) - if _is_aoa_slicemap(A.slicemap) - return parent(A) - else - throw(ArgumentError("flatview for AbstractSlices requires inner dimensions to be first and no dimension reordering, but slicemap is $(A.slicemap)")) - end +function flatview(A::AbstractArray{<:AbstractArray}) + throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) end + + +""" + deepmap(f::Base.Callable, x::Any) + deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray{<:...}}) + +Applies `map` at the deepest possible layer of nested arrays. If `A` is not +a nested array, `deepmap` behaves identical to `Base.map`. +""" +function deepmap end +export deepmap + +deepmap(f::Base.Callable, x::Any) = map(f, x) + +deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = map(X -> deepmap(f, X), A) + + """ nestedview(A::AbstractArray{T,M+N}, M::Integer) nestedview(A::AbstractArray{T,2}) @@ -122,7 +170,6 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} s end -@inline innersize(A::AbstractSlices) = _extract_innerdims(size(parent(A)), getslicemap(A)) @inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = innersize(A)[dim] diff --git a/src/util.jl b/src/util.jl index f0f02bd..396d32f 100644 --- a/src/util.jl +++ b/src/util.jl @@ -47,34 +47,3 @@ Base.@pure require_ndims(A::AbstractArray{T,N}, Val_N::Val{N}) where {T,N} = Base.@pure require_ndims(A::AbstractArray{T,M}, Val_N::Val{N}) where {T,M,N} = throw(ArgumentError("Require an array with $N dimensions")) -@inline @generated function _extract_innerdims(obj::Tuple, slicemap::Tuple{Vararg{Union{Colon,Integer}}}) - # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), - # extract only the elements of obj where the slicemap is a Colon. - expr = Expr(:tuple) - slicepars = slicemap.parameters - for i in 1:length(slicepars) - if slicepars[i] <: Colon - push!(expr.args, :(obj[$i])) - end - end - return expr -end - -@inline @generated function _extract_outerdims(obj::Tuple, slicemap::Tuple{Vararg{Union{Colon,Integer}}}) - # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), - # extract only the elements of obj where the slicemap is a Colon. - expr = Expr(:tuple) - slicepars = slicemap.parameters - for i in 1:length(slicepars) - if slicepars[i] <: Integer - push!(expr.args, :(obj[$i])) - end - end - return expr -end - - -function _is_aoa_slicemap(slicemap::Tuple{Vararg{Union{Colon,Integer}}}) - dims = _oneto_tpl(Val(length(slicemap))) - issorted((_extract_innerdims(dims, slicemap)..., _extract_outerdims(dims, slicemap)...)) -end From d6d11d769f332e03c197149a016c64cb9380b3af Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 09/45] WIP partview --- src/base_slices.jl | 3 ++- src/functions.jl | 45 +++++++++++++++++++++++++++++++++------------ 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/src/base_slices.jl b/src/base_slices.jl index f1e2e3e..c651dd2 100644 --- a/src/base_slices.jl +++ b/src/base_slices.jl @@ -21,7 +21,8 @@ end @inline getpartmode(A::Slices) = BaseSlicing(A.slicemap) -@inline flatview(A::Slices) = parent(A) +@inline unpartview(A::Slices) = parent(A) + @inline stacked(A::Slices) = reshape(parent(A), (length(A), prod(size(parent(A))) ÷ length(A))) diff --git a/src/functions.jl b/src/functions.jl index 7389d5c..11298d6 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -8,7 +8,7 @@ Abstract supertype for array partition modes. Use [`getpartmode`](@ref) to get the partition mode of an array. -See also [`partview`](@ref) and [`flatview`](@ref). +See also [`partview`](@ref) and [`unpartview`](@ref). """ abstract type AbstractPartitionMode end export AbstractPartitionMode @@ -42,7 +42,7 @@ export AbstractSlicingMode Get the partitioning mode of `A`. -`partview(flatview(A), getpartmode(A))` must equal `A`, and should have +`partview(unpartview(A), getpartmode(A))` must equal `A`, and should have the same type as `A` if at all possible. `getpartmode` should be an allocation-free O(1) operation, if at all possible. @@ -66,7 +66,7 @@ return the original unpartition array. `partview` should be an allocation-free O(1) operation, if at all possible. -See also [`flatview`](@ref) and [`getpartmode`](@ref). +See also [`unpartview`](@ref) and [`getpartmode`](@ref). """ function partview end export partview @@ -75,31 +75,52 @@ export partview """ - flatview(A::AbstractArray) - flatview(A::AbstractArray{<:AbstractArray{<:...}}) + unpartview(A::AbstractArray) + unpartview(A::AbstractArray{<:AbstractArray{<:...}}) -View array `A` in a suitable flattened form. +View array `A` in unpartitioned form. -`partview(flatview(A), getpartmode(A))` must equal `A`, and should have +`partview(unpartview(A), getpartmode(A))` must equal `A`, and should have the same type as `A` if at all possible. If `A` is not a nested array return `A` itself. If `A` is a partitioned array, return the original unpartition array. -`flatview` should be an allocation-free O(1) operation, if at all possible. +`unpartview` should be an allocation-free O(1) operation, if at all possible. +""" +function unpartview end +export unpartview + +@inline unpartview(A::AbstractArray) = A + +function unpartview(A::AbstractArray{<:AbstractArray}) + throw(ArgumentError("unpartview not implemented nested arrays of type $(nameof(typeof(A)))")) +end + +""" + flatview(A::AbstractArray) + flatview(A::AbstractArray{<:AbstractArray{<:...}}) + +View array `A` in a flattened form, with inner dimensions first. The shape of +the flattened form will depend on the type of `A`. If the `A` is not a +nested array, the return value is `A` itself. Only specific types of nested +arrays are supported. """ function flatview end export flatview @inline flatview(A::AbstractArray) = A +@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) -function flatview(A::AbstractArray{<:AbstractArray}) - throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) +function flatview(A::AbstractSlices) + if _is_aoa_slicemap(A.slicemap) + return unpartview(A) + else + throw(ArgumentError("flatview for AbstractSlices requires inner dimensions to be first and no dimension reordering, but slicemap is $(A.slicemap)")) + end end - - """ deepmap(f::Base.Callable, x::Any) deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray{<:...}}) From 3edae173b17672fcfda722e14319ab7f4e881b16 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 10/45] Remove _innerlength --- src/array_of_similar_arrays.jl | 5 ----- test/array_of_similar_arrays.jl | 2 +- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 14515a4..23f9d85 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -125,11 +125,6 @@ end end -@inline function _innerlength(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} - prod(innersize(A)) -end - - import Base.== (==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = (A.data == B.data) diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 75a38ea..e547c46 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -330,7 +330,7 @@ using StatsBase: cov2cor @test @inferred(deepmap(f, ASA)).data == ASA.data.*2 @test @inferred(innermap(f, ASA)).data == ASA.data.*2 - @test @inferred(ArraysOfArrays._innerlength(VSV)) == N + @test @inferred(prod(ArraysOfArrays.innersize(VSV))) == N end @testset "map and broadcast" begin From d521a4cd48f9bfd031ccf9035bf14446acd70025 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 11/45] STASH partview --- src/base_slices.jl | 2 +- src/functions.jl | 36 ++++++++++++++++++++++++++++++++---- 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/src/base_slices.jl b/src/base_slices.jl index c651dd2..1b6a8ba 100644 --- a/src/base_slices.jl +++ b/src/base_slices.jl @@ -67,7 +67,7 @@ end end -function is_natural_dimorder_partmode(pmode::BaseSlicing) +function is_memordered_partmode(pmode::BaseSlicing) slicemap = pmode.slicemap dims = _oneto_tpl(Val(length(slicemap))) issorted((_extract_innerdims(dims, slicemap)..., _extract_outerdims(dims, slicemap)...)) diff --git a/src/functions.jl b/src/functions.jl index 11298d6..d85bad2 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -56,6 +56,22 @@ function getpartmode(A::AbstractArray{<:AbstractArray}) end +""" + is_memordered_partmode(pmode::AbstractSlicingMode)::Bool + +Check if `pmode` partitions in memory-order. + +If true, inner arrays are stored contiguously in memory in Julia-native +dimension order, and the same is true for the outer dimensions (no dimention +reordering). + +If true, `flatview` and `unpartview` are equivalent. +""" +function is_memordered_partmode end + +is_memordered_partmode(::Unpartitioned) = true + + """ partview(A::AbstractArray, pmode::AbstractSlicingMode) @@ -86,6 +102,9 @@ the same type as `A` if at all possible. If `A` is not a nested array return `A` itself. If `A` is a partitioned array, return the original unpartition array. +If `is_memordered_partmode(getpartmode(A))` is true, `unpartview(A)` is +equivalent to [`flatview(A)`](@ref). + `unpartview` should be an allocation-free O(1) operation, if at all possible. """ function unpartview end @@ -94,9 +113,10 @@ export unpartview @inline unpartview(A::AbstractArray) = A function unpartview(A::AbstractArray{<:AbstractArray}) - throw(ArgumentError("unpartview not implemented nested arrays of type $(nameof(typeof(A)))")) + throw(ArgumentError("unpartview not implemented for nested arrays of type $(nameof(typeof(A)))")) end + """ flatview(A::AbstractArray) flatview(A::AbstractArray{<:AbstractArray{<:...}}) @@ -105,18 +125,26 @@ View array `A` in a flattened form, with inner dimensions first. The shape of the flattened form will depend on the type of `A`. If the `A` is not a nested array, the return value is `A` itself. Only specific types of nested arrays are supported. + +If `is_memordered_partmode(getpartmode(A))` is true, `flatview(A)` is +equivalent to [`unpartview(A)`](@ref). + +`unpartview` should be an allocation-free O(1) operation, if at all possible. """ function flatview end export flatview @inline flatview(A::AbstractArray) = A -@inline flatview(A::AbstractArray{<:AbstractArray}) = throw(ArgumentError("flatview not implemented nested arrays of type $(nameof(typeof(A)))")) +function flatview(A::AbstractArray{<:AbstractArray}) + throw(ArgumentError("flatview not implemented for nested arrays of type $(nameof(typeof(A)))")) +end function flatview(A::AbstractSlices) - if _is_aoa_slicemap(A.slicemap) + pmode = getpartmode(A) + if is_memordered_partmode(pmode) return unpartview(A) else - throw(ArgumentError("flatview for AbstractSlices requires inner dimensions to be first and no dimension reordering, but slicemap is $(A.slicemap)")) + throw(ArgumentError("flatview required memory-ordered partitioning/slicing, but array has partition mode $pmode")) end end From b045115ca19cb14a3d98bc4166f0749e7626e028 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 12/45] STASH partview --- ext/ArraysOfArraysChainRulesCoreExt.jl | 3 +-- src/array_of_similar_arrays.jl | 14 ++------------ src/functions.jl | 12 ++++++++---- 3 files changed, 11 insertions(+), 18 deletions(-) diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index a643ad3..bc0231c 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -19,7 +19,7 @@ mapthunk(::Type{F}, x::T) where {F,T} = _MappedMaybeThunk{Type{F},T}(F, x) @non_differentiable getpartmode(::Any) @non_differentiable innersize(::Any) - +@non_differentiable is_memordered_partmode(::Any) function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, ::Unpartitioned) @@ -33,7 +33,6 @@ end _unpart_flatview_pullback(ΔΩ) = NoTangent(), ΔΩ - function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, partmode::AbstractSlicingMode) return partview(A, partmode), _partview_pullback end diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 23f9d85..80d8b91 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -107,8 +107,8 @@ Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractAr Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) -Base.parent(A::ArrayOfSimilarArrays) = A.data -Base.stack(A::ArrayOfSimilarArrays) = A.data +unpartview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data +Base.stack(A::ArrayOfSimilarArrays) = unpartview(A) function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} new_ET = Base.promote_op(similar, ET) @@ -130,16 +130,6 @@ import Base.== (A.data == B.data) -""" - flatview(A::ArrayOfSimilarArrays{T,M,N,P})::P - -Returns the array of dimensionality `M + N` wrapped by `A`. The shape of -the result may be freely changed without breaking the inner consistency of -`A`. -""" -flatview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data - - Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] diff --git a/src/functions.jl b/src/functions.jl index d85bad2..04258bf 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -45,7 +45,7 @@ Get the partitioning mode of `A`. `partview(unpartview(A), getpartmode(A))` must equal `A`, and should have the same type as `A` if at all possible. -`getpartmode` should be an allocation-free O(1) operation, if at all possible. +`getpartmode` should be a zero-copy O(1) operation, if at all possible. """ function getpartmode end @@ -80,7 +80,7 @@ View array `A` in partitioned form, as an array of arrays. If `A` is not a nested array return `A` itself. If `A` is a partitioned array, return the original unpartition array. -`partview` should be an allocation-free O(1) operation, if at all possible. +`partview` should be a zero-copy O(1) operation, if at all possible. See also [`unpartview`](@ref) and [`getpartmode`](@ref). """ @@ -105,7 +105,7 @@ return the original unpartition array. If `is_memordered_partmode(getpartmode(A))` is true, `unpartview(A)` is equivalent to [`flatview(A)`](@ref). -`unpartview` should be an allocation-free O(1) operation, if at all possible. +`unpartview` should be a zero-copy O(1) operation, if at all possible. """ function unpartview end export unpartview @@ -129,7 +129,11 @@ arrays are supported. If `is_memordered_partmode(getpartmode(A))` is true, `flatview(A)` is equivalent to [`unpartview(A)`](@ref). -`unpartview` should be an allocation-free O(1) operation, if at all possible. +The result of `flatview(A)` will equal either `stack(A)` +(resp. [`stacked(A)`](@ref)) or `reduce(vcat, A)`, depending on the type of +`A` (sliced-array-like or ragged-array-like). + +`unpartview` should be a zero-copy O(1) operation, if at all possible. """ function flatview end export flatview From 5be1d1b44de96c7600ce8d2c7ea55be4f25ed316 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 13/45] Remove function abstract_nestedarray_type --- docs/src/index.md | 1 - src/functions.jl | 26 -------------------------- test/functions.jl | 8 -------- 3 files changed, 35 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index b4c27dc..a1f3307 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -7,7 +7,6 @@ This package also defines and exports the following new functions applicable to * [`nestedview`](@ref) and [`flatview`](@ref) switch between a flat and a nested view of the same data. * [`innersize`](@ref) returns the size of the elements of an array, provided they all have equal size. * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. -* [`abstract_nestedarray_type`](@ref) returns the type of nested `AbstractArray`s for a given innermost element type with multiple layers of nesting. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. diff --git a/src/functions.jl b/src/functions.jl index 04258bf..3a2668d 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -229,29 +229,3 @@ end @inline innersize(tpl::Tuple{T}) where T = size(only(tpl)) @inline innersize(ref::Ref) = size(only(ref)) - - -""" - abstract_nestedarray_type(T_inner::Type, ::Val{ndims_tuple}) - -Return the type of nested `AbstractArray`s. `T_inner` specifies the element -type of the innermost layer of arrays, `ndims_tuple` specifies the -dimensionality of each nesting layer (outer arrays first). - -If `ndims_tuple` is empty, the returns is the (typically scalar) type -`T_inner` itself. -""" -function abstract_nestedarray_type end -export abstract_nestedarray_type - - -Base.@pure function abstract_nestedarray_type(::Type{T_inner}, outer::Val{ndims_tuple}) where {T_inner,ndims_tuple} - _abstract_nestedarray_type_impl(T_inner, ndims_tuple...) -end - -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}) where {T_inner} = T_inner - -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}, N) where {T_inner} = AbstractArray{T_inner, N} - -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}, N, M, ndims_tuple...) where {T_inner} = - AbstractArray{<:_abstract_nestedarray_type_impl(T_inner, M, ndims_tuple...), N} diff --git a/test/functions.jl b/test/functions.jl index 15d32ce..3da5049 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -48,12 +48,4 @@ using StaticArrays @test @inferred(innersize(Aes1)) == (5, 6, 7) @test @inferred(innersize(Aes2)) == (5, 7, 9) end - - - @testset "abstract_nestedarray_type" begin - @test @inferred(abstract_nestedarray_type(Int, Val(()))) == Int - @test @inferred(abstract_nestedarray_type(Int, Val((2,)))) == AbstractArray{Int, 2} - @test @inferred(abstract_nestedarray_type(Float32, Val((2,3,4)))) == - AbstractArray{<:AbstractArray{<:AbstractArray{Float32, 4}, 3}, 2} - end end From edc981c932703d51fb204dc0f913b4bd4f8f5c1c Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 14/45] Define innersize for non-nested arrays --- src/functions.jl | 4 +++- test/functions.jl | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/src/functions.jl b/src/functions.jl index 3a2668d..7aa524d 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -205,6 +205,8 @@ are not of equal size. function innersize end export innersize +innersize(::AbstractArray) = () + function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} s = if !isempty(A) let sz_A = size(first(A)) @@ -220,7 +222,7 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} end end - s + return s end diff --git a/test/functions.jl b/test/functions.jl index 3da5049..da7f4a1 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -38,6 +38,7 @@ using StaticArrays @testset "innersize" begin + @test @inferred(innersize(rand(3,4,5))) == () @test @inferred(innersize([[1, 2, 3], [4, 5, 6]])) == (3,) @test @inferred(innersize([[]])) == (0,) @test @inferred(innersize([2:5])) == (4,) From 0dc77ea8136a8064814df978c32a4905143f84b3 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 15/45] Don't support innersize for Tuples and Refs Breaking, but likely wasn't used a lot. --- src/functions.jl | 3 --- test/functions.jl | 2 -- 2 files changed, 5 deletions(-) diff --git a/src/functions.jl b/src/functions.jl index 7aa524d..bd09208 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -228,6 +228,3 @@ end @inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = innersize(A)[dim] - -@inline innersize(tpl::Tuple{T}) where T = size(only(tpl)) -@inline innersize(ref::Ref) = size(only(ref)) diff --git a/test/functions.jl b/test/functions.jl index da7f4a1..ea7dbe8 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -42,8 +42,6 @@ using StaticArrays @test @inferred(innersize([[1, 2, 3], [4, 5, 6]])) == (3,) @test @inferred(innersize([[]])) == (0,) @test @inferred(innersize([2:5])) == (4,) - @test @inferred(innersize((2:5,))) == (4,) - @test @inferred(innersize(Ref(2:5))) == (4,) @test_throws DimensionMismatch @inferred(innersize([[1, 2, 3], [4, 5]])) @test @inferred(innersize(Aes1)) == (5, 6, 7) From 449cdabfa0c52cf42202ce522dcd5fff6f9e6da0 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 16/45] STASH splitview --- ext/ArraysOfArraysChainRulesCoreExt.jl | 33 +-- src/array_of_similar_arrays.jl | 32 ++- src/base_slices.jl | 44 ++-- src/functions.jl | 294 +++++++++++++++++-------- src/util.jl | 9 - test/Project.toml | 1 + test/base_slices.jl | 38 ++++ test/functions.jl | 97 +++++--- test/runtests.jl | 2 + test/static_arrays.jl | 17 ++ test/testdefs.jl | 170 ++++++++++++++ 11 files changed, 563 insertions(+), 174 deletions(-) create mode 100644 test/base_slices.jl create mode 100644 test/static_arrays.jl create mode 100644 test/testdefs.jl diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index bc0231c..9f70020 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -2,10 +2,11 @@ module ArraysOfArraysChainRulesCoreExt -using ChainRulesCore: ChainRulesCore, NoTangent, Thunk, unthunk, @thunk, @non_differentiable +using ChainRulesCore: ChainRulesCore, NoTangent, AbstractThunk, Thunk, unthunk, @thunk, @non_differentiable +using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitview, fused, flatview, innersize +using ArraysOfArrays: NonSplitMode, AbstractSlicingMode using ArraysOfArrays: ArrayOfSimilarArrays -using ArraysOfArrays: flatview struct _MappedMaybeThunk{F, T} <: AbstractThunk @@ -17,32 +18,32 @@ mapthunk(f::F, x::T) where {F,T} = _MappedMaybeThunk{F,T}(f, x) mapthunk(::Type{F}, x::T) where {F,T} = _MappedMaybeThunk{Type{F},T}(F, x) -@non_differentiable getpartmode(::Any) +@non_differentiable getsplitmode(::Any) @non_differentiable innersize(::Any) -@non_differentiable is_memordered_partmode(::Any) +@non_differentiable is_memordered_splitmode(::Any) -function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, ::Unpartitioned) - return flatview(A), _partview_pullback +function ChainRulesCore.rrule(::typeof(splitview), A::AbstractArray, ::NonSplitMode) + return fused(A), _partview_pullback end _unpart_partview_pullback(ΔΩ) = NoTangent(), ΔΩ, NoTangent() -function ChainRulesCore.rrule(::typeof(flatview), A::AbstractArray) - return flatview(A), _unpart_flatview_pullback +function ChainRulesCore.rrule(::typeof(fused), A::AbstractArray) + return fused(A), _unpart_joinedview_pullback end -_unpart_flatview_pullback(ΔΩ) = NoTangent(), ΔΩ +_unpart_joinedview_pullback(ΔΩ) = NoTangent(), ΔΩ -function ChainRulesCore.rrule(::typeof(partview), A::AbstractArray, partmode::AbstractSlicingMode) - return partview(A, partmode), _partview_pullback +function ChainRulesCore.rrule(::typeof(splitview), A::AbstractArray, smode::AbstractSlicingMode) + return splitview(A, smode), _partview_pullback end -_partview_pullback(ΔΩ) = NoTangent(), mapthunk(flatview, ΔΩ), NoTangent() +_partview_pullback(ΔΩ) = NoTangent(), mapthunk(fused, ΔΩ), NoTangent() -function ChainRulesCore.rrule(::typeof(flatview), A::AbstractArray{<:AbstractArray}) - pmode = getpartmode(A) - return flatview(A), Base.Fix2(_flatview_pullback, pmode) +function ChainRulesCore.rrule(::typeof(fused), A::AbstractArray{<:AbstractArray}) + smode = getsplitmode(A) + return fused(A), Base.Fix2(_joinedview_pullback, smode) end -_flatview_pullback(ΔΩ, pmode) = NoTangent(), mapthunk(Base.Fix2(partview, pmode), ΔΩ) +_joinedview_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitview, smode), ΔΩ) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 80d8b91..0bdc900 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -107,8 +107,8 @@ Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractAr Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) -unpartview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data -Base.stack(A::ArrayOfSimilarArrays) = unpartview(A) +fused(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data +Base.stack(A::ArrayOfSimilarArrays) = fused(A) function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} new_ET = Base.promote_op(similar, ET) @@ -359,3 +359,31 @@ Statistics.cov(X::AbstractVectorOfSimilarVectors; corrected::Bool = true) = Statistics.cor(X::AbstractVectorOfSimilarVectors) = cor(flatview(X); dims = 2) + + +""" + nestedview(A::AbstractArray{T,M+N}, M::Integer) + nestedview(A::AbstractArray{T,2}) + +AbstractArray{<:AbstractArray{T,M},N} + +View array `A` in as an `N`-dimensional array of `M`-dimensional arrays by +wrapping it into an [`ArrayOfSimilarArrays`](@ref). + +It's also possible to use a `StaticVector` of length `S` as the type of the +inner arrays via + + nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S}}) + nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S,T}}) +""" +function nestedview end +export nestedview + +@inline nestedview(A::AbstractArray{T,L}, M::Integer) where {T,L} = + ArrayOfSimilarArrays{T,M}(A) + +@inline nestedview(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = + ArrayOfSimilarArrays{T,M}(A) + +@inline nestedview(A::AbstractArray{T,2}) where {T} = + VectorOfSimilarVectors{T}(A) diff --git a/src/base_slices.jl b/src/base_slices.jl index 1b6a8ba..6fd3aa6 100644 --- a/src/base_slices.jl +++ b/src/base_slices.jl @@ -3,7 +3,7 @@ """ struct BaseSlicing{N,TPL<:Tuple{Vararg{Int,N}}} <: AbstractSlicingMode -The partitioning mode of `Slices`. +The split mode of `Slices`. Constructor: @@ -15,23 +15,20 @@ BaseSlicing(slicemap::Tuple{Vararg{Union{Colon,Int}}}}) See also [`AbstractSlicingMode`](@ref). """ -struct BaseSlicing{N,TPL<:Tuple{Vararg{Union{Colon,Int},N}}} <: AbstractSlicingMode +struct BaseSlicing{M,N,TPL<:Tuple{Vararg{Union{Colon,Int}}}} <: AbstractSlicingMode{M,N} slicemap::TPL end +export BaseSlicing +is_memordered_splitmode(::BaseSlicing{1,1,Tuple{Colon,Int}}) = true -@inline getpartmode(A::Slices) = BaseSlicing(A.slicemap) -@inline unpartview(A::Slices) = parent(A) - - -@inline stacked(A::Slices) = reshape(parent(A), (length(A), prod(size(parent(A))) ÷ length(A))) - -@inline innersize(A::AbstractSlices) = getinnnerdims(size(flatview(A)), getpartmode(A)) +function is_memordered_splitmode(smode::BaseSlicing{M,N}) where {M,N} + dims = _oneto_tpl(Val(M+N)) + issorted((getinnerdims(dims, smode)..., getouterdims(dims, smode)...)) +end -@inline @generated function getinnnerdims(obj::Tuple, pmode::BaseSlicing{N,SliceMapT}) where {N,SliceMapT} - # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), - # extract only the elements of obj where the slicemap is a Colon. +@inline @generated function getinnerdims(obj::Tuple, ::BaseSlicing{M,N,SliceMapT}) where {M,N,SliceMapT} expr = Expr(:tuple) slicepars = SliceMapT.parameters for i in 1:length(slicepars) @@ -43,9 +40,7 @@ end end -@inline @generated function getouterdims(obj::Tuple, pmode::BaseSlicing{N,SliceMapT}) where {N,SliceMapT} - # slicemap may be something like (Colon(), 2, Colon(), 1, Colon()), - # extract only the elements of obj where the slicemap is a Colon. +@inline @generated function getouterdims(obj::Tuple, smode::BaseSlicing{M,N,SliceMapT}) where {M,N,SliceMapT} slicepars = SliceMapT.parameters outdimidxs = Int[] for i in 1:length(slicepars) @@ -59,7 +54,7 @@ end result_expr = Expr(:tuple, [:(obj[outdimidxs[idxorder[$i]]]) for i in eachindex(outdimidxs)]...) quote - slicemap = pmode.slicemap + slicemap = smode.slicemap outdimidxs = $outdimidxs_expr idxorder = $idxorder_expr return $result_expr @@ -67,8 +62,17 @@ end end -function is_memordered_partmode(pmode::BaseSlicing) - slicemap = pmode.slicemap - dims = _oneto_tpl(Val(length(slicemap))) - issorted((_extract_innerdims(dims, slicemap)..., _extract_outerdims(dims, slicemap)...)) +@inline function getsplitmode(A::Slices) + M = ndims(eltype(A)) + N = ndims(A) + slicemap = A.slicemap + BaseSlicing{M,N,typeof(slicemap)}(slicemap) end + +function splitview(A::AbstractArray, smode::BaseSlicing) + slicemap = smode.slicemap + axs = getouterdims(axes(A), smode) + return Slices(A, slicemap, axs) +end + +@inline fused(A::Slices) = parent(A) diff --git a/src/functions.jl b/src/functions.jl index bd09208..ee40470 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -2,138 +2,199 @@ """ - abstract type AbstractPartitionMode + abstract type AbstractSplitMode -Abstract supertype for array partition modes. +Abstract supertype for array split modes. -Use [`getpartmode`](@ref) to get the partition mode of an array. +Use [`getsplitmode`](@ref) to get the split mode of an array. -See also [`partview`](@ref) and [`unpartview`](@ref). +See also [`splitview`](@ref) and [`fused`](@ref). """ -abstract type AbstractPartitionMode end -export AbstractPartitionMode +abstract type AbstractSplitMode end +export AbstractSplitMode + +""" + struct NonSplitMode{N} <: AbstractSplitMode + +The split mode of unsplit collections that have `N` dimensions. + +Constructor: `NonSplitMode{N}()` """ - struct Unpartitioned <: AbstractPartitionMode +struct NonSplitMode{N} <: AbstractSplitMode end +export NonSplitMode + + +""" + struct UnknownSplitMode{AT} <: AbstractSplitMode + +Split mode of generic split objects of type{T} that have been split in an +unknown way, e.g. nested arrays of type `Array{<:Array}`. -The partitioning mode of unpartitioned arrays. +Since the split parts may be (typically are) non-contiguous in memory, this +split mode not allow for `unsplitview` or `flatview`. It is also not +inferrable of the split object should be interpeted as a sliced array, +a ragged array, or something else. -Constructor: `Unpartitioned()` +Constructor: `UnknownSplitMode{T}()` """ -struct Unpartitioned <: AbstractPartitionMode end -export Unpartitioned +struct UnknownSplitMode{AT} <: AbstractSplitMode end +export UnknownSplitMode """ - abstract type AbstractSlicingMode <: AbstractPartitionMode + abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode -Abstract supertype for array partition modes. +Abstract supertype for array slicing modes with `M` inner dimensions and `N` +outer dimensions. -Use `getpartmode` to get the partition mode of an partitioned array. +Use `getsplitmode` to get the split mode of an split array. """ -abstract type AbstractSlicingMode <: AbstractPartitionMode end +abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode end export AbstractSlicingMode """ - getpartmode(A::AbstractArray)::Unpartitioned - getpartmode(A::AbstractArray{<:AbstractArray})::AbstractSlicingMode + getsplitmode(A::AbstractArray)::NonSplitMode + getsplitmode(A::AbstractArray{<:AbstractArray})::AbstractSplitMode -Get the partitioning mode of `A`. +Get the split mode of `A`. -`partview(unpartview(A), getpartmode(A))` must equal `A`, and should have -the same type as `A` if at all possible. +`splitview(fused(A), getsplitmode(A))` must equal `A`, and should have +the same type as `A` if at all possible, except if `getsplitmode(A)` is an +`UnknownSplitMode`. -`getpartmode` should be a zero-copy O(1) operation, if at all possible. +`getsplitmode` should be a zero-copy O(1) operation, if at all possible. """ -function getpartmode end +function getsplitmode end +export getsplitmode -@inline getpartmode(::AbstractArray) = Unpartitioned() +@inline getsplitmode(::T) where T = UnknownSplitMode{T}() -function getpartmode(A::AbstractArray{<:AbstractArray}) - throw(ArgumentError("getpartmode not implemented for nested arrays of type $(nameof(typeof(A)))")) -end +@inline getsplitmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() + +@inline getsplitmode(A::AbstractArray{<:AbstractArray}) = UnknownSplitMode{typeof(A)}() """ - is_memordered_partmode(pmode::AbstractSlicingMode)::Bool + is_memordered_splitmode(smode::AbstractSplitMode)::Bool -Check if `pmode` partitions in memory-order. +Check if `smode` splits in memory-order. If true, inner arrays are stored contiguously in memory in Julia-native dimension order, and the same is true for the outer dimensions (no dimention reordering). -If true, `flatview` and `unpartview` are equivalent. +If true, `flatview` and `fused` are equivalent. """ -function is_memordered_partmode end +function is_memordered_splitmode end +export is_memordered_splitmode -is_memordered_partmode(::Unpartitioned) = true +is_memordered_splitmode(::NonSplitMode) = true +is_memordered_splitmode(::UnknownSplitMode) = false """ - partview(A::AbstractArray, pmode::AbstractSlicingMode) + ArraysOfArrays.getinnerdims(tpl::Tuple, smode::AbstractSlicing) -View array `A` in partitioned form, as an array of arrays. +Get the entries of `tpl` corresponding to the inner dimensions of slicing +mode `smode`, in the order specified by `smode`. +""" +function getinnerdims end -If `A` is not a nested array return `A` itself. If `A` is a partitioned array, -return the original unpartition array. +@inline getinnerdims(::Tuple, ::NonSplitMode) = () + +function getinnerdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getinnerdims cannot be used with UnknownSplitMode")) +end -`partview` should be a zero-copy O(1) operation, if at all possible. -See also [`unpartview`](@ref) and [`getpartmode`](@ref). """ -function partview end -export partview + ArraysOfArrays.getouterdims(tpl::Tuple, smode::AbstractSlicing) -@inline partview(A::AbstractArray, ::Unpartitioned) = A +Get the entries of `tpl` corresponding to the outer dimensions of slicing +mode `smode`, in the order specified by `smode`. +""" +function getouterdims end + +@inline getouterdims(x::Tuple, ::NonSplitMode) = x + +function getouterdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getouterdims cannot be used with UnknownSplitMode")) +end + + +""" + splitview(A::AbstractArray, smode::AbstractSplitMode) + +View array `A` in split form, as an array of arrays. + +`splitview` should be a zero-copy O(1) operation, if at all possible. + +See also [`fused`](@ref) and [`getsplitmode`](@ref). +""" +function splitview end +export splitview + +@inline splitview(obj::Any, ::NonSplitMode) = obj + +function splitview(::Any, ::UnknownSplitMode) + throw(ArgumentError("splitview cannot be used with UnknownSplitMode")) +end """ - unpartview(A::AbstractArray) - unpartview(A::AbstractArray{<:AbstractArray{<:...}}) + fused(A::AbstractArray) + fused(A::AbstractArray{<:AbstractArray}) -View array `A` in unpartitioned form. +View array `A` in unsplit form. -`partview(unpartview(A), getpartmode(A))` must equal `A`, and should have -the same type as `A` if at all possible. +`splitview(fused(A), getsplitmode(A))` must equal `A`, and should have +the same type as `A` if at all possible, except if `getsplitmode(A)` is an +`UnknownSplitMode`. -If `A` is not a nested array return `A` itself. If `A` is a partitioned array, -return the original unpartition array. +If `A` is not a nested array return `A` itself. If `A` is a split array, +return the original unsplit array. -If `is_memordered_partmode(getpartmode(A))` is true, `unpartview(A)` is +If `is_memordered_splitmode(getsplitmode(A))` is true, `fused(A)` is equivalent to [`flatview(A)`](@ref). -`unpartview` should be a zero-copy O(1) operation, if at all possible. +`fused` should be a zero-copy O(1) operation, if at all possible. """ -function unpartview end -export unpartview +function fused end +export fused + +@inline fused(obj) = _joinedview_impl(obj, getsplitmode(obj)) + +@inline fused(A::AbstractArray) = A -@inline unpartview(A::AbstractArray) = A +@inline fused(A::AbstractArray{<:AbstractArray}) = _joinedview_impl(A, getsplitmode(A)) -function unpartview(A::AbstractArray{<:AbstractArray}) - throw(ArgumentError("unpartview not implemented for nested arrays of type $(nameof(typeof(A)))")) +@inline _joinedview_impl(obj, ::NonSplitMode) = obj + +function _joinedview_impl(@nospecialize(obj), ::UnknownSplitMode) + throw(ArgumentError("fused not implemented for objects of type $(nameof(typeof(obj))) with unknown split mode")) end """ flatview(A::AbstractArray) - flatview(A::AbstractArray{<:AbstractArray{<:...}}) + flatview(A::AbstractArray{<:AbstractArray}) View array `A` in a flattened form, with inner dimensions first. The shape of the flattened form will depend on the type of `A`. If the `A` is not a nested array, the return value is `A` itself. Only specific types of nested arrays are supported. -If `is_memordered_partmode(getpartmode(A))` is true, `flatview(A)` is -equivalent to [`unpartview(A)`](@ref). +If `is_memordered_splitmode(getsplitmode(A))` is true, `flatview(A)` is +equivalent to [`fused(A)`](@ref). The result of `flatview(A)` will equal either `stack(A)` (resp. [`stacked(A)`](@ref)) or `reduce(vcat, A)`, depending on the type of `A` (sliced-array-like or ragged-array-like). -`unpartview` should be a zero-copy O(1) operation, if at all possible. +`flatview` should be a zero-copy O(1) operation, if at all possible. """ function flatview end export flatview @@ -144,60 +205,53 @@ function flatview(A::AbstractArray{<:AbstractArray}) end function flatview(A::AbstractSlices) - pmode = getpartmode(A) - if is_memordered_partmode(pmode) - return unpartview(A) + smode = getsplitmode(A) + if is_memordered_splitmode(smode) + return fused(A) else - throw(ArgumentError("flatview required memory-ordered partitioning/slicing, but array has partition mode $pmode")) + throw(ArgumentError("flatview required memory-ordered split/slicing, but array has split mode $smode")) end end - -""" - deepmap(f::Base.Callable, x::Any) - deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray{<:...}}) - -Applies `map` at the deepest possible layer of nested arrays. If `A` is not -a nested array, `deepmap` behaves identical to `Base.map`. """ -function deepmap end -export deepmap - -deepmap(f::Base.Callable, x::Any) = map(f, x) - -deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = map(X -> deepmap(f, X), A) + stacked(A::AbstractArray) + stacked(A::AbstractArray{<:AbstractArray}) +Join stacked arrays of a nested array into a single array along one or more +new dimensions, return non-nested arrays unchanged. +Similar to `Base.stack`, but can return the original underlying array of +sliced arrays in more cases. """ - nestedview(A::AbstractArray{T,M+N}, M::Integer) - nestedview(A::AbstractArray{T,2}) +function stacked end +export stacked -AbstractArray{<:AbstractArray{T,M},N} +@inline stacked(A::AbstractArray) = A +@inline stacked(A::AbstractArray{<:AbstractArray}) = _stacked_impl(A, getsplitmode(A)) -View array `A` in as an `N`-dimensional array of `M`-dimensional arrays by -wrapping it into an [`ArrayOfSimilarArrays`](@ref). +_stacked_impl(A::AbstractArray{<:AbstractArray}, ::AbstractSplitMode) = stack(A) -It's also possible to use a `StaticVector` of length `S` as the type of the -inner arrays via +function _stacked_impl(A::AbstractSlices, smode::AbstractSlicingMode) + A_joined = fused(A) + is_memordered_splitmode(smode) ? A_joined : _stacked_permutedims(A_joined, smode) +end - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S}}) - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S,T}}) -""" -function nestedview end -export nestedview +function _stacked_permutedims(A_joined::AbstractArray{T,N}, smode::AbstractSlicingMode) where {T,N} + dimnumbers = _oneto_tpl(Val(N)) + dimorder = (getinnerdims(dimnumbers, smode)..., getouterdims(_dimnumbers, smode)...) + return permutedims(A_joined, dimorder)::typeof(A_joined) +end -@inline nestedview(A::AbstractArray{T,L}, M::Integer) where {T,L} = - ArrayOfSimilarArrays{T,M}(A) -@inline nestedview(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = - ArrayOfSimilarArrays{T,M}(A) +# ToDo: Add function `unstacked(A::AbstractArray, smode::AbstractSlicingMode)`? -@inline nestedview(A::AbstractArray{T,2}) where {T} = - VectorOfSimilarVectors{T}(A) +# ToDo: Add function `stacked_splitmode(smode::AbstractSlicingMode)`, and +# `stacked_splitmode(A::AbstractArray{<:AbstractArray}) = stacked_splitmode(getsplitmode(A))`, +# so that `splitview(stacked(A), stacked_splitmode(A) == A`? """ - innersize(A:AbstractArray{<:AbstractArray}, [dim]) + innersize(A::AbstractArray{<:AbstractArray}, [dim]) Returns the size of the element arrays of `A`. Fails if the element arrays are not of equal size. @@ -205,7 +259,9 @@ are not of equal size. function innersize end export innersize -innersize(::AbstractArray) = () +@inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = innersize(A)[dim] + +innersize(::AbstractArray{<:Number}) = () function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} s = if !isempty(A) @@ -225,6 +281,50 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} return s end +@inline innersize(A::AbstractSlices) = getinnerdims(size(fused(A)), getsplitmode(A)) + -@inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = - innersize(A)[dim] +""" + innermap(f, A::AbstractArray) + innermap(f, A::AbstractArray{<:AbstractArray}) + +Nested `map` at depth 2. Equivalent to `map(X -> map(f, X) A)` for arrays +of arrays, otherwise equivalent to `Base.map`. +""" +function innermap end +export innermap + +innermap(f, obj) = _generic_innermap_impl(f, obj) +innermap(f, A::AbstractArray) = map(f, A) +innermap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(map, f), A) +innermap(f, A::AbstractSlices) = _generic_innermap_impl(f, A) + +function _generic_innermap_impl(f, obj) + joined_obj = fused(obj) + mapped_joined_obj = map(f, joined_obj) + mapped_obj = splitview(mapped_joined_obj, getsplitmode(obj)) + return mapped_obj +end + + +""" + deepmap(f, A::AbstractArray) + deepmap(f, A::AbstractArray{<:AbstractArray{<:...}}) + +Applies `map` at the deepest layer of nested arrays. If `A` is not +a nested array, `deepmap` behaves identical to `Base.map`. +""" +function deepmap end +export deepmap + +deepmap(f, obj) = _generic_deepmap_impl(f, obj) +deepmap(f, A::AbstractArray) = map(f, A) +deepmap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(deepmap, f), A) +deepmap(f, A::AbstractSlices) = _generic_deepmap_impl(f, A) + +function _generic_deepmap_impl(f, obj) + joined_obj = fused(obj) + mapped_joined_obj = deepmap(f, joined_obj) + mapped_obj = splitview(mapped_joined_obj, getsplitmode(obj)) + return mapped_obj +end diff --git a/src/util.jl b/src/util.jl index 396d32f..5e8b68b 100644 --- a/src/util.jl +++ b/src/util.jl @@ -1,15 +1,6 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). -#= - -function _split_dims(dims::NTuple{N,Integer}) where {N} - int_dims = Int.(dims) - Base.front(int_dims), int_dims[end] -end - -=# - @inline _tail_impl(x, ys...) = (ys...,) @inline _tail(x) = _tail_impl(x...) diff --git a/test/Project.toml b/test/Project.toml index 225e375..ea0ca2f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" diff --git a/test/base_slices.jl b/test/base_slices.jl new file mode 100644 index 0000000..3abd618 --- /dev/null +++ b/test/base_slices.jl @@ -0,0 +1,38 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using ArraysOfArrays: getinnerdims, getouterdims + +include("testdefs.jl") + +@testset "base_slices" begin + A_orig = rand(5,6,7,8,9) + A_orig_mat = rand(5,6) + + Aes1 = eachslice(A_orig; dims = (4,5)) + @test @inferred(getsplitmode(Aes1)) isa BaseSlicing{3,2,Tuple{Colon,Colon,Colon,Int,Int}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes1))) == (1,2,3) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes1))) == (4,5) + + Aes2 = eachslice(A_orig; dims = (3,1,5)) + @test @inferred(getsplitmode(Aes2)) isa BaseSlicing{2,3,Tuple{Int,Colon,Int,Colon,Int}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes2))) == (2,4) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes2))) == (3,1,5) + + Aec = eachcol(A_orig_mat) + @test @inferred(getsplitmode(Aec)) isa BaseSlicing{1,1,Tuple{Colon,Int}} + @test @inferred(getinnerdims((1,2), getsplitmode(Aec))) == (1,) + @test @inferred(getouterdims((1,2), getsplitmode(Aec))) == (2,) + + Aer = eachrow(A_orig_mat) + @test @inferred(getsplitmode(Aer)) isa BaseSlicing{1,1,Tuple{Int,Colon}} + @test @inferred(getinnerdims((1,2), getsplitmode(Aer))) == (2,) + @test @inferred(getouterdims((1,2), getsplitmode(Aer))) == (1,) + + test_api(Aes1, Array(Aes1), A_orig) + test_api(Aes2, Array(Aes2), A_orig) + test_api(Aec, Array(Aec), A_orig_mat) + test_api(Aer, Array(Aer), A_orig_mat) +end diff --git a/test/functions.jl b/test/functions.jl index ea7dbe8..7a5af8f 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -3,48 +3,85 @@ using ArraysOfArrays using Test -using StaticArrays +using ArraysOfArrays: getinnerdims, getouterdims +include("testdefs.jl") @testset "functions" begin - Aes1 = eachslice(rand(5,6,7,8,9); dims = (4,5)) - Aes2 = eachslice(rand(5,6,7,8,9); dims = (4,2)) - - function gen_nested() - A11 = [1 2; 3 4] - A21 = [4 5 6; 7 8 9] - A12 = [10 11; 12 13; 14 15] - A22 = [16 17; 18 19] - hcat([A11, A21], [A12, A22]) - end + A_0 = fill(Float32(42)) + A_0_flat = fill(Float32(42)) + + A_1 = Float32[3, 8, 4, 6] + A_1_flat = Float32[3, 8, 4, 6] + + A_1e = Float32[] + A_1e_flat = Float32[] + + A_2 = [Float32[3 8; 4 6], Float32[1 2; 5 7; 9 0]] + A_2_flat = Float32[3, 8, 4, 6, 1, 2, 5, 7, 9, 0] + + A_2e = Matrix{Float32}[] + A_2e_flat = Matrix{Float32}(undef, 0, 0) + A_2b = [Float32[3 8; 4 6], Float32[1 2; 9 0]] + A_2b_flat = stack(A_2b) - @testset "flatview and nestedview" begin - A = [(@SArray randn(3, 2, 4)) for i in 1:2, j in 1:2] - @test @inferred(nestedview(flatview(A), Val(3))) == A + A_3 = [[Float32[3, 8, 4], Float32[1, 2, 5]], [Float32[6, 7], Float32[9, 0]]] + A_3_flat = reduce(vcat, A_3) - B = rand(3, 2, 4) - @test @inferred(nestedview(flatview(B), SVector{3})) == @inferred(nestedview(B, Val(1))) + @testset "getsplitmode" begin + @test @inferred(getsplitmode(A_0)) isa NonSplitMode{0} + @test @inferred(getsplitmode(A_1)) isa NonSplitMode{1} + @test @inferred(getsplitmode(A_1e)) isa NonSplitMode{1} + @test @inferred(getsplitmode(A_2)) isa UnknownSplitMode{typeof(A_2)} + @test @inferred(getsplitmode(A_2e)) isa UnknownSplitMode{typeof(A_2e)} + @test @inferred(getsplitmode(A_2b)) isa UnknownSplitMode{typeof(A_2b)} + @test @inferred(getsplitmode(A_3)) isa UnknownSplitMode{typeof(A_3)} + end - @test @inferred flatview(Aes1) == parent(Aes1) - @test_throws ArgumentError @inferred flatview(Aes2) + @testset "innersize" begin + @test @inferred(innersize(A_0)) == () + @test @inferred(innersize(A_1)) == () + @test @inferred(innersize(A_1e)) == () + @test_throws DimensionMismatch innersize(A_2) + @test @inferred(innersize(A_2e)) == (0, 0) + @test @inferred(innersize(A_2b)) == (2, 2) + @test @inferred(innersize(A_3)) == (2,) end + @testset "innermap" begin + f = x -> x^2 + @test @inferred(innermap(f, A_0)) == fill(1764) + @test @inferred(innermap(f, A_1)) == Float32[9, 64, 16, 36] + @test @inferred(innermap(f, A_2)) == [Float32[9 64; 16 36], Float32[1 4; 25 49; 81 0]] + @test_throws MethodError innermap(f, A_3) - @testset "getslicemap" begin - @test @inferred(getslicemap(Aes1)) == (:, :, :, 1, 2) - @test @inferred(getslicemap(Aes2)) == (:, 2, :, 1, :) + @test @inferred(innermap(length, A_0)) == fill(1) + @test @inferred(innermap(length, A_1)) == fill(1, 4) + @test @inferred(innermap(length, A_2)) == [fill(1, 2, 2), fill(1, 3, 2)] + @test @inferred(innermap(length, A_3)) == [fill(3, 2), fill(2, 2)] end + @testset "deepmap" begin + f = x -> x^2 + @test @inferred(deepmap(f, A_0)) == fill(1764) + @test @inferred(deepmap(f, A_1)) == Float32[9, 64, 16, 36] + @test @inferred(deepmap(f, A_2)) == [Float32[9 64; 16 36], Float32[1 4; 25 49; 81 0]] + @test @inferred(deepmap(f, A_3)) == [[Float32[9, 64, 16], Float32[1, 4, 25]], [Float32[36, 49], Float32[81, 0]]] - @testset "innersize" begin - @test @inferred(innersize(rand(3,4,5))) == () - @test @inferred(innersize([[1, 2, 3], [4, 5, 6]])) == (3,) - @test @inferred(innersize([[]])) == (0,) - @test @inferred(innersize([2:5])) == (4,) - @test_throws DimensionMismatch @inferred(innersize([[1, 2, 3], [4, 5]])) - - @test @inferred(innersize(Aes1)) == (5, 6, 7) - @test @inferred(innersize(Aes2)) == (5, 7, 9) + @test @inferred(deepmap(length, A_0)) == fill(1) + @test @inferred(deepmap(length, A_1)) == fill(1, 4) + @test @inferred(deepmap(length, A_2)) == [fill(1, 2, 2), fill(1, 3, 2)] + @test @inferred(deepmap(length, A_3)) == [[fill(1, 3), fill(1, 3)], [fill(1, 2), fill(1, 2)]] + end + + @testset "Nested array API" begin + test_api(A_0, A_0, A_0_flat) + test_api(A_1, A_1, A_1_flat) + test_api(A_1e, A_1e, A_1e_flat) + test_api(A_2, A_2, A_2_flat) + test_api(A_2e, A_2e, A_2e_flat) + test_api(A_2b, A_2b, A_2b_flat) + test_api(A_3, A_3, A_3_flat) end end diff --git a/test/runtests.jl b/test/runtests.jl index fa08bb9..cabda5d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,8 +5,10 @@ import Test Test.@testset "Package ArraysOfArrays" begin include("test_aqua.jl") include("functions.jl") + include("base_slices.jl") include("array_of_similar_arrays.jl") include("vector_of_arrays.jl") include("broadcasting.jl") + include("static_arrays.jl") include("test_docs.jl") end # testset diff --git a/test/static_arrays.jl b/test/static_arrays.jl new file mode 100644 index 0000000..ee136a4 --- /dev/null +++ b/test/static_arrays.jl @@ -0,0 +1,17 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using StaticArrays + + +@testset "StaticArray Extension" begin + @testset "flatview and nestedview" begin + A = [(@SArray randn(3, 2, 4)) for i in 1:2, j in 1:2] + @test @inferred(nestedview(flatview(A), Val(3))) == A + + B = rand(3, 2, 4) + @test @inferred(nestedview(flatview(B), SVector{3})) == @inferred(nestedview(B, Val(1))) + end +end diff --git a/test/testdefs.jl b/test/testdefs.jl new file mode 100644 index 0000000..69ab712 --- /dev/null +++ b/test/testdefs.jl @@ -0,0 +1,170 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent + +#if !isdefined(Main, :test_api) + maptest_f(x::Number) = x^2 + maptest_f(x::AbstractArray{<:Number}) = sum(x)^2 + maptest_f(x::AbstractArray) = length(x)^2 + + function test_notangent_rrule(f, args::Vararg{Any,N}) where {N} + y, f_pullback = @inferred(rrule(f, args...)) + @test y == f(args...) + dy = NoTangent() + @test @inferred(f_pullback(dy)) == ntuple(_ -> NoTangent(), Val(N+1)) + end + + function test_joinedview_rrule(A) + @testset "Test fused rrule for $(nameof(typeof(A)))" begin + A_joined, f_pullback = @inferred(rrule(fused, A)) + @test A_joined == fused(A) + dy = Thunk(Returns(A_joined)) + tangents = @inferred(f_pullback(dy)) + @test tangents isa Tuple{NoTangent,<:AbstractThunk} + @test unthunk(tangents[2]) == A + end + end + + function test_splitview_rrule(A, smode) + @testset "Test splitview rrule for $(nameof(typeof(A)))" begin + A_split, f_pullback = @inferred(rrule(splitview, A, smode)) + @test A_split == splitview(A, smode) + dy = Thunk(Returns(A_split)) + tangents = @inferred(f_pullback(dy)) + @test tangents isa Tuple{NoTangent,<:AbstractThunk,NoTangent} + @test unthunk(tangents[2]) == A + end + end + + function test_api(A, A_array_ref, A_unsplit_ref) + @testset "Test API for $(nameof(typeof(A)))" begin + global g_state = (;A, A_array_ref, A_unsplit_ref) + + @test Array(A) isa Array{<:Any,ndims(A)} + A_array = Array(A) + @test A == A_array + @test isequal(A, A_array) + @test isequal(A_array, A_array_ref) + + @test @inferred(getsplitmode(A)) isa AbstractSplitMode + smode = getsplitmode(A) + test_notangent_rrule(getsplitmode, A) + + @test @inferred(is_memordered_splitmode(smode)) isa Bool + test_notangent_rrule(is_memordered_splitmode, smode) + + if A isa AbstractSlices + let M = ndims(eltype(A)), N = ndims(A) + @test smode isa AbstractSlicingMode{M,N} + end + elseif eltype(A) <: Number + @test smode isa NonSplitMode{ndims(A)} + end + + @test @inferred(eltype(A)) <: Union{Number,AbstractArray} + T_elem = eltype(A) + + innersz = if smode isa NonSplitMode + @inferred(innersize(A)) + elseif smode isa AbstractSlicingMode + @inferred(innersize(A)) + else + let sz = (try innersize(A); catch err; err; end) + sz isa Exception ? sz : @inferred(innersize(A)) + end + end + + if !(innersz isa Exception) + test_notangent_rrule(innersize, A) + end + + @test innersz isa Union{Exception, Dims} + + if !isempty(A) + A_elem_1 = @inferred(A[first(eachindex(A))]) + @test typeof(A_elem_1) == T_elem + if !(innersz isa Exception) + @test all(size.(A) .== Ref(innersz)) + end + + @inferred(innermap(maptest_f, A)) == innermap(maptest_f, Array(A)) + @inferred(deepmap(maptest_f, A)) == deepmap(maptest_f, Array(A)) + end + + _smode_M(::AbstractSlicingMode{M,N}) where {M,N} = M + _smode_N(::AbstractSlicingMode{M,N}) where {M,N} = N + + dimstpl = ntuple(identity, Val(ndims(A_unsplit_ref))) + + if smode isa AbstractSlicingMode + M, N = _smode_M(smode), _smode_N(smode) + A_array_stacked = stack(A_array) + @test M == ndims(eltype(A)) + @test N == ndims(A) + + @test Array(stack(A)) == A_array_stacked + + if is_memordered_splitmode(smode) + if A isa Slices + # stack(A) never returns parent for Slices, even if possible: + @test @inferred(stack(A)) == A_unsplit_ref + else + @test @inferred(stack(A)) === A_unsplit_ref + end + @test @inferred(stacked(A)) === A_unsplit_ref + @test @inferred(flatview(A)) === A_unsplit_ref + else + @test Array(@inferred(stack(A))) == A_array_stacked + @test Array(@inferred(stacked(A))) == A_array_stacked + @test_throws ArgumentError flatview(A) + end + + @test @inferred(getinnerdims(dimstpl, smode)) isa NTuple{M,Int} + @test @inferred(getouterdims(dimstpl, smode)) isa NTuple{N,Int} + innerdims = getinnerdims(dimstpl, smode) + outerdims = getouterdims(dimstpl, smode) + @test permutedims(A_unsplit_ref, (innerdims..., outerdims...)) == A_array_stacked + elseif smode isa NonSplitMode + @test @inferred(stacked(A)) === A + else + stacked_A = try stack(A); catch err; err; end + if stacked_A isa Exception + @test_throws typeof(stacked_A) stacked(A) + else + @inferred(stacked(A)) == stacked_A + end + end + + if smode isa UnknownSplitMode + @test @inferred(is_memordered_splitmode(smode)) == false + @test_throws ArgumentError fused(A) + @test_throws ArgumentError flatview(A) + @test_throws ArgumentError splitview(A_unsplit_ref, smode) + @test_throws ArgumentError getinnerdims(dimstpl, smode) + @test_throws ArgumentError getouterdims(dimstpl, smode) + elseif smode isa NonSplitMode + @test @inferred(is_memordered_splitmode(smode)) == true + @test @inferred(fused(A)) === A + @test @inferred(flatview(A)) === A + @test @inferred(splitview(A, smode)) === A + @test @inferred(getinnerdims(dimstpl, smode)) == () + @test @inferred(getouterdims(dimstpl, smode)) == dimstpl + + test_joinedview_rrule(A) + test_splitview_rrule(A, smode) + else + if A isa Slices + @test fused(A) === parent(A) + end + @test @inferred(fused(A)) == A_unsplit_ref + A_unsplit = fused(A) + @test typeof(A_unsplit) == typeof(A_unsplit_ref) + @test typeof(splitview(A_unsplit, smode)) == typeof(A) + @test splitview(A_unsplit, smode) == A + + test_joinedview_rrule(A) + test_splitview_rrule(A_unsplit, smode) + end + end + end +#end From fccf08ccca1f486ec014306107205cdbe20f4a06 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 17/45] Rename ArrayOfSimilarArrays to SlicedView --- README.md | 4 +- docs/src/index.md | 6 +- ext/ArraysOfArraysAdaptExt.jl | 6 +- ext/ArraysOfArraysChainRulesCoreExt.jl | 16 +-- src/ArraysOfArrays.jl | 2 +- src/array_of_similar_arrays.jl | 136 ++++++++++++++----------- test/array_of_similar_arrays.jl | 66 ++++++------ test/broadcasting.jl | 6 +- 8 files changed, 128 insertions(+), 114 deletions(-) diff --git a/README.md b/README.md index efd24e9..57b2c1b 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,8 @@ A Julia package for efficient storage and handling of nested arrays. -ArraysOfArrays provides two different types of nested arrays: `ArrayOfSimilarArrays` and `VectorOfArrays`. -An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. +ArraysOfArrays provides two different types of nested arrays: `SlicedView` and `VectorOfArrays`. +An `SlicedView` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. ## Documentation diff --git a/docs/src/index.md b/docs/src/index.md index a1f3307..4acdbf6 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,6 +1,6 @@ # ArraysOfArrays.jl -A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: [`ArrayOfSimilarArrays`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). +A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: [`SlicedView`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). This package also defines and exports the following new functions applicable to nested arrays in general: @@ -10,9 +10,9 @@ This package also defines and exports the following new functions applicable to * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. -## [ArrayOfSimilarArrays](@id section_ArrayOfSimilarArrays) +## [SlicedView](@id section_ArrayOfSimilarArrays) -An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays: +An `SlicedView` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays: ```julia A_flat = rand(2,3,4,5,6) diff --git a/ext/ArraysOfArraysAdaptExt.jl b/ext/ArraysOfArraysAdaptExt.jl index 6bd2e65..6ae9d2d 100644 --- a/ext/ArraysOfArraysAdaptExt.jl +++ b/ext/ArraysOfArraysAdaptExt.jl @@ -5,13 +5,13 @@ module ArraysOfArraysAdaptExt import Adapt using Adapt: adapt -using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays +using ArraysOfArrays: SlicedView, VectorOfArrays using ArraysOfArrays: no_consistency_checks -function Adapt.adapt_structure(to, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +function Adapt.adapt_structure(to, A::SlicedView{T,M,N}) where {T,M,N} adapted_data = adapt(to, A.data) - ArrayOfSimilarArrays{eltype(adapted_data),M,N}(adapted_data) + SlicedView{eltype(adapted_data),M,N}(adapted_data) end diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index 9f70020..3312754 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -6,7 +6,7 @@ using ChainRulesCore: ChainRulesCore, NoTangent, AbstractThunk, Thunk, unthunk, using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitview, fused, flatview, innersize using ArraysOfArrays: NonSplitMode, AbstractSlicingMode -using ArraysOfArrays: ArrayOfSimilarArrays +using ArraysOfArrays: SlicedView struct _MappedMaybeThunk{F, T} <: AbstractThunk @@ -48,24 +48,24 @@ _joinedview_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitview, s function _aosa_ctor_fromflat_pullback(ΔΩ) - NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) + NoTangent(), flatview(convert(SlicedView, unthunk(ΔΩ))) end -function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} - return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback +function ChainRulesCore.rrule(::Type{SlicedView{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} + return SlicedView{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback end _aosa_ctor_fromnested_pullback(ΔΩ) = NoTangent(), ΔΩ -function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} - return ArrayOfSimilarArrays{T,M,N}(A), _aosa_ctor_fromnested_pullback +function ChainRulesCore.rrule(::Type{SlicedView{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} + return SlicedView{T,M,N}(A), _aosa_ctor_fromnested_pullback end -function ChainRulesCore.rrule(::typeof(flatview), A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +function ChainRulesCore.rrule(::typeof(flatview), A::SlicedView{T,M,N}) where {T,M,N} function flatview_pullback(ΔΩ) data = unthunk(ΔΩ) - NoTangent(), ArrayOfSimilarArrays{eltype(data),M,N}(data) + NoTangent(), SlicedView{eltype(data),M,N}(data) end return flatview(A), flatview_pullback diff --git a/src/ArraysOfArrays.jl b/src/ArraysOfArrays.jl index 2a0788d..06973dd 100644 --- a/src/ArraysOfArrays.jl +++ b/src/ArraysOfArrays.jl @@ -8,7 +8,7 @@ __precompile__(true) Efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: -[`ArrayOfSimilarArrays`](@ref section_ArrayOfSimilarArrays) and +[`SlicedView`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). """ module ArraysOfArrays diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 0bdc900..5fd92c4 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -33,7 +33,7 @@ export AbstractVectorOfSimilarVectors """ - ArrayOfSimilarArrays{T,M,N,P} <: AbstractArrayOfSimilarArrays{T,M,N} + SlicedView{T,M,N,P} <: AbstractArrayOfSimilarArrays{T,M,N} Represents a view of an array of dimension `M + N` as an array of dimension M with elements that are arrays with dimension N. All element arrays @@ -41,8 +41,8 @@ implicitly have equal size/axes. Constructors: - ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray) - ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray) + SlicedView{T,M,N}(flat_data::AbstractArray) + SlicedView{T,M}(flat_data::AbstractArray) The following type aliases are defined: @@ -64,14 +64,14 @@ A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T flatview(A_nested) === A_flat ``` """ -struct ArrayOfSimilarArrays{ +struct SlicedView{ T, M, N, P<:AbstractArray{T}, ET<:AbstractArray{T,M} } <: AbstractArrayOfSimilarArrays{T,M,N,ET} data::P - function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U}) where {T,M,N,U} + function SlicedView{T,M,N}(flat_data::AbstractArray{U}) where {T,M,N,U} require_ndims(flat_data, _add_vals(Val{M}(), Val{N}())) conv_parent = _convert_elype(T, flat_data) P = typeof(conv_parent) @@ -80,135 +80,135 @@ struct ArrayOfSimilarArrays{ end end -function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U}) where {T,M,U} +function SlicedView{T,M}(flat_data::AbstractArray{U}) where {T,M,U} _, size_outer = split_tuple(size(flat_data), Val{M}()) N = length(size_outer) - ArrayOfSimilarArrays{T,M,N}(flat_data) + SlicedView{T,M,N}(flat_data) end -export ArrayOfSimilarArrays +export SlicedView -function ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} - B = ArrayOfSimilarArrays{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...)) +function SlicedView{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} + B = SlicedView{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...)) copyto!(B, A) end -ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - ArrayOfSimilarArrays{T,M,N}(A) +SlicedView{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = + SlicedView{T,M,N}(A) -ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = - ArrayOfSimilarArrays{T,M,N}(A) +SlicedView(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = + SlicedView{T,M,N}(A) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M}}, flat_data::AbstractArray{U}) where {T,M,U} = R(flat_data) +Base.convert(R::Type{SlicedView{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} = R(flat_data) +Base.convert(R::Type{SlicedView{T,M}}, flat_data::AbstractArray{U}) where {T,M,U} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) +Base.convert(R::Type{SlicedView{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) +Base.convert(R::Type{SlicedView{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) +Base.convert(R::Type{SlicedView}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) -fused(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data -Base.stack(A::ArrayOfSimilarArrays) = fused(A) +fused(A::SlicedView{T,M,N}) where {T,M,N} = A.data +Base.stack(A::SlicedView) = fused(A) -function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} +function Base.Array(A::SlicedView{T,M,N,P,ET}) where {T,M,N,P,ET} new_ET = Base.promote_op(similar, ET) return Array{new_ET,N}(A) end -function getslicemap(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +function getslicemap(::SlicedView{T,M,N}) where {T,M,N} return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) end #!!!!! -@inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +@inline function innersize(A::SlicedView{T,M,N}) where {T,M,N} front_tuple(size(A.data), Val{M}()) end import Base.== -(==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = +(==)(A::SlicedView{T,M,N}, B::SlicedView{T,M,N}) where {T,M,N} = (A.data == B.data) -Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] +Base.size(A::SlicedView{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] -Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} = +Base.@propagate_inbounds Base.getindex(A::SlicedView{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} = view(A.data, _ncolons(Val{M}())..., idxs...) -Base.@propagate_inbounds Base.setindex!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = +Base.@propagate_inbounds Base.setindex!(A::SlicedView{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = setindex!(A.data, x, _ncolons(Val{M}())..., idxs...) -Base.@propagate_inbounds function Base.unsafe_view(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} +Base.@propagate_inbounds function Base.unsafe_view(A::SlicedView{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} dataview = view(A.data, _ncolons(Val{M}())..., idxs...) L = length(size(dataview)) N_view = L - M - ArrayOfSimilarArrays{T,M,N_view}(dataview) + SlicedView{T,M,N_view}(dataview) end -@inline function Base.resize!(A::ArrayOfSimilarArrays{T,M,N}, dims::Vararg{Integer,N}) where {T,M,N} +@inline function Base.resize!(A::SlicedView{T,M,N}, dims::Vararg{Integer,N}) where {T,M,N} resize!(A.data, innersize(A)..., dims...) A end -function Base.similar(A::ArrayOfSimilarArrays{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U} +function Base.similar(A::SlicedView{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U} data = A.data size_inner, size_outer = split_tuple(size(data), Val{M}()) # ToDo: Don't use similar if data is an ElasticArray? - ArrayOfSimilarArrays{T,M,N}(similar(data, U, size_inner..., dims...)) + SlicedView{T,M,N}(similar(data, U, size_inner..., dims...)) end -function Base.deepcopy(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} - ArrayOfSimilarArrays{T,M,N}(deepcopy(A.data)) +function Base.deepcopy(A::SlicedView{T,M,N}) where {T,M,N} + SlicedView{T,M,N}(deepcopy(A.data)) end -function Base.copyto!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} +function Base.copyto!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} copyto!(dest.data, src.data) dest end -function Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} +function Base.append!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't append, shape of element arrays of source and dest are not equal")) append!(dest.data, src.data) dest end -Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - append!(dest, ArrayOfSimilarArrays(src)) +Base.append!(dest::SlicedView{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = + append!(dest, SlicedView(src)) -function Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} +function Base.prepend!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't prepend, shape of element arrays of source and dest are not equal")) prepend!(dest.data, src.data) dest end -Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - prepend!(dest, ArrayOfSimilarArrays(src)) +Base.prepend!(dest::SlicedView{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = + prepend!(dest, SlicedView(src)) -function innermap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +function innermap(f::Base.Callable, A::SlicedView{T,M,N}) where {T,M,N} new_data = map(f, A.data) U = eltype(new_data) - ArrayOfSimilarArrays{U,M,N}(new_data) + SlicedView{U,M,N}(new_data) end -function deepmap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} +function deepmap(f::Base.Callable, A::SlicedView{T,M,N}) where {T,M,N} new_data = deepmap(f, A.data) U = eltype(new_data) - ArrayOfSimilarArrays{U,M,N}(new_data) + SlicedView{U,M,N}(new_data) end -Base.map(::typeof(identity), A::ArrayOfSimilarArrays) = A -Base.Broadcast.broadcasted(::typeof(identity), A::ArrayOfSimilarArrays) = A +Base.map(::typeof(identity), A::SlicedView) = A +Base.Broadcast.broadcasted(::typeof(identity), A::SlicedView) = A Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) = @@ -221,15 +221,15 @@ const VectorOfSimilarArrays{ T, M, P<:AbstractArray{T}, ET<:AbstractArray{T} -} = ArrayOfSimilarArrays{T,M,1,P,ET} +} = SlicedView{T,M,1,P,ET} export VectorOfSimilarArrays VectorOfSimilarArrays{T}(flat_data::AbstractArray{U}) where {T,U} = - ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) + SlicedView{T,length(Base.front(size(flat_data))),1}(flat_data) VectorOfSimilarArrays(flat_data::AbstractArray{T}) where {T} = - ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) + SlicedView{T,length(Base.front(size(flat_data))),1}(flat_data) VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = VectorOfSimilarArrays{T,M}(A) @@ -267,7 +267,7 @@ function Base.pushfirst!(V::VectorOfSimilarArrays{T,M}, x::AbstractArray{U,M}) w end # Will need equivalent of resize! that resizes in front of data instead of in back: -# popfirst!(V::ArrayOfSimilarArrays) = ... +# popfirst!(V::SlicedView) = ... function _empty_data_size(A::VectorOfSimilarArrays{T,M}) where {T,M} @@ -293,15 +293,15 @@ const ArrayOfSimilarVectors{ T, N, P<:AbstractArray{T}, ET<:AbstractVector{T} -} = ArrayOfSimilarArrays{T,1,N,P,ET} +} = SlicedView{T,1,N,P,ET} export ArrayOfSimilarVectors ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U}) where {T,U} = - ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) + SlicedView{T,1,length(Base.front(size(flat_data)))}(flat_data) ArrayOfSimilarVectors(flat_data::AbstractArray{T}) where {T} = - ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) + SlicedView{T,1,length(Base.front(size(flat_data)))}(flat_data) ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = ArrayOfSimilarVectors{T,N}(A) @@ -320,18 +320,18 @@ const VectorOfSimilarVectors{ T, P<:AbstractArray{T,2}, ET<:AbstractVector{T} -} = ArrayOfSimilarArrays{T,1,1,P,ET} +} = SlicedView{T,1,1,P,ET} export VectorOfSimilarVectors VectorOfSimilarVectors{T}(flat_data::AbstractArray{U,2}) where {T,U} = - ArrayOfSimilarArrays{T,1,1}(flat_data) + SlicedView{T,1,1}(flat_data) VectorOfSimilarVectors(flat_data::AbstractArray{T,2}) where {T} = VectorOfSimilarVectors{T}(flat_data) VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} = - ArrayOfSimilarArrays{T,1}(A) + SlicedView{T,1}(A) VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} = VectorOfSimilarVectors{T}(A) @@ -368,7 +368,7 @@ Statistics.cor(X::AbstractVectorOfSimilarVectors) = AbstractArray{<:AbstractArray{T,M},N} View array `A` in as an `N`-dimensional array of `M`-dimensional arrays by -wrapping it into an [`ArrayOfSimilarArrays`](@ref). +wrapping it into an [`SlicedView`](@ref). It's also possible to use a `StaticVector` of length `S` as the type of the inner arrays via @@ -380,10 +380,24 @@ function nestedview end export nestedview @inline nestedview(A::AbstractArray{T,L}, M::Integer) where {T,L} = - ArrayOfSimilarArrays{T,M}(A) + SlicedView{T,M}(A) @inline nestedview(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = - ArrayOfSimilarArrays{T,M}(A) + SlicedView{T,M}(A) @inline nestedview(A::AbstractArray{T,2}) where {T} = VectorOfSimilarVectors{T}(A) + + +# Deprecated: +const ArrayOfSimilarArrays{T,M,N} = SlicedView{T,M,N} +export ArrayOfSimilarArrays + +const VectorOfSimilarArrays{T,M} = VoASlicedView{T,M} +export VectorOfSimilarArrays + +const ArrayOfSimilarVectors{T,N} = AoVSlicedView{T,N} +export ArrayOfSimilarVectors + +const VectorOfSimilarVectors{T} = VoVSlicedView{T} +export VectorOfSimilarVectors diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index e547c46..0c0e173 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -32,8 +32,8 @@ using StatsBase: cov2cor @inline function test_from_flat(::Type{TT}, ::Type{RT}, Val_L::Val{L}) where { - TT<:ArrayOfSimilarArrays, - RT<:ArrayOfSimilarArrays, + TT<:SlicedView, + RT<:SlicedView, L } @testset "$TT from Array{Float64,$L}" begin @@ -56,8 +56,8 @@ using StatsBase: cov2cor @inline function test_from_nested(::Type{TT}, ::Type{RT}, Val_M::Val{M}, Val_N::Val{N}) where { - TT<:ArrayOfSimilarArrays, - RT<:ArrayOfSimilarArrays, + TT<:SlicedView, + RT<:SlicedView, M, N } @testset "$TT from Array{Array{Float64,$M},$N}" begin @@ -84,33 +84,33 @@ using StatsBase: cov2cor @testset "construct/convert from flat array" begin - test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(SlicedView{Float64,2,3}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(SlicedView{Float64,2}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(5)) test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) - test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(SlicedView{Float32,2,3}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(SlicedView{Float32,2}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(5)) test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) - test_rrule(ArrayOfSimilarArrays{Float64,2,2}, rand(2,3,4,5)) + test_rrule(SlicedView{Float64,2,2}, rand(2,3,4,5)) end @testset "construct/convert from nested arrays" begin - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) + test_from_nested(SlicedView, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(SlicedView{Float64,2,3}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(SlicedView{Float64}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(SlicedView{Float32}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) + test_from_nested(SlicedView, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(SlicedView, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(SlicedView, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) @@ -129,12 +129,12 @@ using StatsBase: cov2cor r = @inferred(rand(5,5)) @test @inferred(flatview(ArrayOfSimilarVectors(r))) == r - test_rrule(ArrayOfSimilarArrays{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) + test_rrule(SlicedView{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) end @testset "AbstractSlices interface" begin - @test @inferred(ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices - let A = ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7)) + @test @inferred(SlicedView{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices + let A = SlicedView{Float32,2,3}(rand(3,4,5,6,7)) @test @inferred(A[2, 3, 4]) isa AbstractArray{Float32,2} ref_ET = typeof(A[2, 3, 4]) @test A isa AbstractSlices{ref_ET, 3} @@ -150,8 +150,8 @@ using StatsBase: cov2cor EA_ref2 = rand_flat_array(Val(3)) EA1 = ElasticArray{Float64, 3}(EA_ref1) EA2 = ElasticArray{Float64, 3}(EA_ref2) - AEA1 = ArrayOfSimilarArrays{Float64, 3}(EA1) - AEA2 = ArrayOfSimilarArrays{Float64, 3}(EA2) + AEA1 = SlicedView{Float64, 3}(EA1) + AEA2 = SlicedView{Float64, 3}(EA2) AEA1_ref = copy(AEA1) AEA2_ref = copy(AEA2) append!(AEA1, AEA2) @@ -161,8 +161,8 @@ using StatsBase: cov2cor @test @inferred(reverse(AEA2.data, dims=3)[:,:,i]) == @inferred(AEA1.data[:,:,N+1-i]) end - A1 = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) - A2 = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) + A1 = SlicedView{Float64,1}(rand_flat_array(Val(1))) + A2 = SlicedView{Float64,1}(rand_flat_array(Val(1))) A1_data = copy(A1.data) A2_data = copy(A2.data) append!(A1, A2) @@ -173,14 +173,14 @@ using StatsBase: cov2cor end @testset "similar and copyto!" begin - A = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) + A = SlicedView{Float64,1}(rand_flat_array(Val(1))) @test (@inferred copyto!((@inferred similar(A)), A)) == A - A = ArrayOfSimilarArrays{Float64,2}(rand_flat_array(Val(5))) + A = SlicedView{Float64,2}(rand_flat_array(Val(5))) @test (@inferred copyto!((@inferred similar(A)), A)) == A A_data = rand_flat_array(Val(4)) - A = ArrayOfSimilarArrays{Float64, 2}(A_data) + A = SlicedView{Float64, 2}(A_data) A_similar = similar(A, Array{Float64, 2}, size(A)) @test @inferred(size(A)) == @inferred(size(A_similar)) @test @inferred(size(A.data)) == @inferred(size(A_similar.data)) @@ -198,7 +198,7 @@ using StatsBase: cov2cor @testset "deepcopy" begin - A = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) + A = SlicedView{Float64,1}(rand_flat_array(Val(1))) @test (@inferred deepcopy(A)) == A @test typeof(deepcopy(A)) == typeof(A) end @@ -206,14 +206,14 @@ using StatsBase: cov2cor @testset "flatview" begin A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) + B = SlicedView(A) @inferred(flatview(B))[:] == collect(flatview(A)) end @testset "empty" begin A = [rand(2,3), rand(2,3), rand(2,3)] - B = ArrayOfSimilarArrays(A) + B = SlicedView(A) @test typeof(@inferred empty(B)) == typeof(B) @test empty(A) == empty(B) @@ -227,10 +227,10 @@ using StatsBase: cov2cor @testset "stats" begin VV = [rand(3) for i in 1:10] - VV_aosa = ArrayOfSimilarArrays(VV) + VV_aosa = SlicedView(VV) VA = [rand(2,3,3) for i in 1:10] - VA_aosa = ArrayOfSimilarArrays(VA) + VA_aosa = SlicedView(VA) array_cmp(A, B) = (A ≈ B) && (size(A) == size(B)) @@ -279,7 +279,7 @@ using StatsBase: cov2cor # ------------------------------------------------------------------- A_flat = rand(2,3,4,5) - ASA = @inferred(ArrayOfSimilarArrays{Float64,2,2}(A_flat)) + ASA = @inferred(SlicedView{Float64,2,2}(A_flat)) @test ASA.data == A_flat # ------------------------------------------------------------------- @@ -317,7 +317,7 @@ using StatsBase: cov2cor r = vcat(r1,r2,r3,r4) VSV = VectorOfSimilarVectors(r) VSA = VectorOfSimilarArrays(r) - ASA = ArrayOfSimilarArrays([r1,r2,r3,r4]) + ASA = SlicedView([r1,r2,r3,r4]) f = x -> x.*2 diff --git a/test/broadcasting.jl b/test/broadcasting.jl index a2823b3..a9d889a 100644 --- a/test/broadcasting.jl +++ b/test/broadcasting.jl @@ -27,7 +27,7 @@ using Test end end - let A = ArrayOfSimilarArrays(ref_AosA1(Float32, 100)) + let A = SlicedView(ref_AosA1(Float32, 100)) refA = Array(A) for Idxs in [ @@ -45,7 +45,7 @@ using Test ] refA = Array(A) - @test @inferred(broadcast(getindex, A, Idxs...)) isa ArrayOfSimilarArrays{eltype(eltype(A))} + @test @inferred(broadcast(getindex, A, Idxs...)) isa SlicedView{eltype(eltype(A))} @test getindex.(A, Idxs...) == getindex.(refA, Idxs...) end end @@ -54,7 +54,7 @@ using Test @testset "findall" begin for A in [ VectorOfArrays(ref_VoA1(Bool, 100)), - ArrayOfSimilarArrays(ref_AosA1(Bool, 100)) + SlicedView(ref_AosA1(Bool, 100)) ] refA = Array(A) From a0e6211a3823bfb0b127a291ab0ad1b44f2507a8 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 18/45] Rename VectorOfArrays to PartsView --- docs/src/index.md | 2 +- src/broadcasting.jl | 4 ++-- src/vector_of_arrays.jl | 48 ++++++++++++++++++++++++---------------- test/broadcasting.jl | 4 ++-- test/vector_of_arrays.jl | 32 +++++++++++++-------------- 5 files changed, 50 insertions(+), 40 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 4acdbf6..fb718ed 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -106,7 +106,7 @@ all(x -> x == 4.2, VA[2]) ### Type aliases The following type aliases are defined: -* `VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +* `PartsView{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` ### Appending data and resizing diff --git a/src/broadcasting.jl b/src/broadcasting.jl index f253fb5..e76cad2 100644 --- a/src/broadcasting.jl +++ b/src/broadcasting.jl @@ -45,7 +45,7 @@ end # ToDo: Extend to vectors of arrays. function Base.Broadcast.broadcasted( ::typeof(getindex), - A::Union{VectorOfSimilarVectors, VectorOfVectors}, + A::Union{VectorOfSimilarVectors, PartsView}, Idxs::Union{AbstractVector{<:AbstractVector{<:Integer}},AbstractVector{Colon},_RefLike{<:Union{AbstractVector{<:Integer},Colon}}}... ) # Checks size compatibility: @@ -123,7 +123,7 @@ end function Base.Broadcast.broadcasted( ::typeof(findall), - A::Union{VectorOfSimilarVectors{Bool}, VectorOfVectors{Bool}} + A::Union{VectorOfSimilarVectors{Bool}, PartsView{Bool}} ) new_lengths = _similar_idx_vector(A, _idx_type(A), length(A)) new_lengths .= sum.(A) diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 645a9e2..791d61c 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -35,10 +35,10 @@ VectorOfArrays( Other suitable values for `checks` are `ArraysOfArrays.simple_consistency_checks` and `ArraysOfArrays.no_consistency_checks`. -`VectorOfVectors` is defined as an type alias: +`PartsView` is defined as an type alias: ```julia -`VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +`PartsView{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` ``` """ struct VectorOfArrays{ @@ -421,15 +421,15 @@ Base.Broadcast.broadcasted(::typeof(identity), A::VectorOfArrays) = A """ - VectorOfVectors{T,...} = VectorOfArrays{T,1,...} + PartsView{T,...} = VectorOfArrays{T,1,...} Constructors: ```julia -VectorOfVectors(A::AbstractVector{<:AbstractVector}) -VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} +PartsView(A::AbstractVector{<:AbstractVector}) +PartsView{T}(A::AbstractVector{<:AbstractVector}) where {T} -VectorOfVectors( +PartsView( data::AbstractVector, elem_ptr::AbstractVector{<:Integer}, checks::Function = full_consistency_checks ) @@ -437,21 +437,21 @@ VectorOfVectors( See also [VectorOfArrays](@ref). ``` """ -const VectorOfVectors{ +const PartsView{ T, VT<:AbstractVector{T}, VI<:AbstractVector{<:Integer}, VD<:AbstractVector{Dims{0}} } = VectorOfArrays{T,1,0,VT,VI,VD} -export VectorOfVectors +export PartsView -VectorOfVectors{T}() where {T} = VectorOfArrays{T,1}() +PartsView{T}() where {T} = VectorOfArrays{T,1}() -VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A) -VectorOfVectors(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A) +PartsView{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A) +PartsView(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A) -VectorOfVectors( +PartsView( data::AbstractVector, elem_ptr::AbstractVector{I}, checks::Function = full_consistency_checks @@ -468,7 +468,7 @@ VectorOfVectors( consgrouped_ptrs(A::AbstractVector) Compute an element pointer vector, suitable for creation of a -`VectorOfVectors` that implies grouping equal consecutive entries of +`PartsView` that implies grouping equal consecutive entries of `A`. Example: @@ -476,7 +476,7 @@ Example: ```julia A = [1, 1, 2, 3, 3, 2, 2, 2] elem_ptr = consgrouped_ptrs(A) - first.(VectorOfVectors(A, elem_ptr)) == [1, 2, 3, 2] + first.(PartsView(A, elem_ptr)) == [1, 2, 3, 2] ``` consgrouped_ptrs Typically, `elem_ptr` will be used to apply the computed grouping to other @@ -484,7 +484,7 @@ data: ```julia B = [1, 2, 3, 4, 5, 6, 7, 8] - VectorOfVectors(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]] + PartsView(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]] ``` """ function consgrouped_ptrs end @@ -516,7 +516,7 @@ end Compute a grouping of equal consecutive elements on `source` via [`consgrouped_ptrs`](@ref) and apply the grouping to target, resp. each element of `target`. `target` may be an vector or a named or unnamed tuple of -vectors. The result is a `VectorOfVectors`, resp. a tuple of such. +vectors. The result is a `PartsView`, resp. a tuple of such. Example: @@ -561,15 +561,25 @@ export consgroupedview function consgroupedview(source::AbstractVector, target::AbstractVector) elem_ptr = consgrouped_ptrs(source) - VectorOfVectors(target, elem_ptr) + PartsView(target, elem_ptr) end function consgroupedview(source::AbstractVector, target::NTuple{N,AbstractVector}) where {N} elem_ptr = consgrouped_ptrs(source) - map(X -> VectorOfVectors(X, elem_ptr), target) + map(X -> PartsView(X, elem_ptr), target) end function consgroupedview(source::AbstractVector, target::NamedTuple{syms,<:NTuple{N,AbstractVector}}) where {syms,N} elem_ptr = consgrouped_ptrs(source) - map(X -> VectorOfVectors(X, elem_ptr), target) + map(X -> PartsView(X, elem_ptr), target) end + + +# Deprecated: + +const VectorOfArrays{T,N} = PartsView{T,N} +export VectorOfArrays + + +const VectorOfVectors{T} = VoVPartsView{T} +export VectorOfVectors diff --git a/test/broadcasting.jl b/test/broadcasting.jl index a9d889a..cacf3c9 100644 --- a/test/broadcasting.jl +++ b/test/broadcasting.jl @@ -18,7 +18,7 @@ using Test for Idxs in [ ([rand(eachindex(a), rand(1:length(a))) for a in A],), - (VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),), + (PartsView([rand(eachindex(a), rand(1:length(a))) for a in A]),), (tuple(1:1),), (tuple([1, 1, 1]),), (tuple(:),), (Ref(1:1),), (Ref([1, 1, 1]),), (Ref(:),), ] @@ -32,7 +32,7 @@ using Test for Idxs in [ ([rand(eachindex(a), rand(1:length(a))) for a in A],), - (VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),), + (PartsView([rand(eachindex(a), rand(1:length(a))) for a in A]),), ] @test @inferred(broadcast(getindex, A, Idxs...)) isa VectorOfArrays{eltype(eltype(A))} @test getindex.(A, Idxs...) == getindex.(refA, Idxs...) diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 8ea4d00..6021347 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -32,10 +32,10 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test @inferred(VectorOfArrays{Float64,1}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} @test VectorOfArrays{Float64,1}(deepcopy(A1)) == A1 - @test @inferred(VectorOfVectors(deepcopy(A1))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors(deepcopy(A1)) == A1 - @test @inferred(VectorOfVectors{Float64}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors{Float64}(deepcopy(A1)) == A1 + @test @inferred(PartsView(deepcopy(A1))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView(deepcopy(A1)) == A1 + @test @inferred(PartsView{Float64}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView{Float64}(deepcopy(A1)) == A1 A1_empty = ref_AoA1(Float32, 0) @test @inferred(VectorOfArrays(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} @@ -45,10 +45,10 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test @inferred(VectorOfArrays{Float64,1}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} @test VectorOfArrays{Float64,1}(deepcopy(A1_empty)) == A1_empty - @test @inferred(VectorOfVectors(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors(deepcopy(A1_empty)) == A1_empty - @test @inferred(VectorOfVectors{Float64}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors{Float64}(deepcopy(A1_empty)) == A1_empty + @test @inferred(PartsView(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView(deepcopy(A1_empty)) == A1_empty + @test @inferred(PartsView{Float64}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView{Float64}(deepcopy(A1_empty)) == A1_empty A2 = ref_AoA2(Float32, 4) @test @inferred(VectorOfArrays(deepcopy(A2))) isa VectorOfArrays{Float32,2,1,Array{Float32,1},Array{Int,1},Array{Tuple{Int},1}} @@ -128,7 +128,7 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test getindex_of_UR == V12 @test getindex_of_vector == getindex_of_UR - VV = @inferred(VectorOfVectors{Float64}()) + VV = @inferred(PartsView{Float64}()) data = @inferred(rand(5)) @inferred(push!(VV, data)) @test @inferred(getindex(VV, 1)) == data @@ -253,7 +253,7 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr # ------------------------------------------------------------------- - VV = @inferred(VectorOfVectors{Float64}()) + VV = @inferred(PartsView{Float64}()) d1 = @inferred(rand(5)) d2 = @inferred(rand(4)) @@ -290,21 +290,21 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr A_grouped_ref = [[1, 1], [2], [3, 3], [2, 2, 2]] elem_ptr = consgrouped_ptrs(A) elem_ptr32 = Int32.(consgrouped_ptrs(A)) - @test first.(@inferred(VectorOfVectors(A, elem_ptr))) == [1, 2, 3, 2] - @test first.(@inferred(VectorOfVectors(A, elem_ptr32))) == [1, 2, 3, 2] + @test first.(@inferred(PartsView(A, elem_ptr))) == [1, 2, 3, 2] + @test first.(@inferred(PartsView(A, elem_ptr32))) == [1, 2, 3, 2] B = [1, 2, 3, 4, 5, 6, 7, 8] B_grouped_ref = [[1, 2], [3], [4, 5], [6, 7, 8]] - @test @inferred(VectorOfVectors(B, elem_ptr)) == B_grouped_ref - @test @inferred(VectorOfVectors(B, elem_ptr32)) == B_grouped_ref + @test @inferred(PartsView(B, elem_ptr)) == B_grouped_ref + @test @inferred(PartsView(B, elem_ptr32)) == B_grouped_ref C = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8] C_grouped_ref = [[1.1, 2.2], [3.3], [4.4, 5.5], [6.6, 7.7, 8.8]] - @test @inferred(consgroupedview(A, B)) isa VectorOfVectors + @test @inferred(consgroupedview(A, B)) isa PartsView @test consgroupedview(A, B) == B_grouped_ref - @test @inferred(consgroupedview(A, (B, C))) isa NTuple{2, VectorOfVectors} + @test @inferred(consgroupedview(A, (B, C))) isa NTuple{2, PartsView} @test consgroupedview(A, (B, C)) == (B_grouped_ref, C_grouped_ref) data = (a = A, b = B, c = C) From cd810de820b8a15aefbbea5b4a0cf0b527b79387 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 19/45] Make AbstractSplitMode a Function --- src/functions.jl | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/src/functions.jl b/src/functions.jl index ee40470..fccfac5 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -2,17 +2,36 @@ """ - abstract type AbstractSplitMode + abstract type AbstractSplitMode <: Function Abstract supertype for array split modes. Use [`getsplitmode`](@ref) to get the split mode of an array. +Use [`splitview`](@ref) or call `smode::AbstractSplitMode` as a function to +split an array: + +```julia +splitview(A, smode) === smode(A) +``` + See also [`splitview`](@ref) and [`fused`](@ref). + +# Implementation + +Subtypes of `AbstractSplitMode` should specialize + +* `splitview(A, smode::SomeSplitMode)` for arrays `A` +* `is_memordered_splitmode(smode::SomeSplitMode)` + +`(smode::SomeSplitMode)(A)` calls `splitview(A, smode)` by default and should +not be specialized. """ -abstract type AbstractSplitMode end +abstract type AbstractSplitMode <: Function end export AbstractSplitMode +(smode::AbstractSplitMode)(A::AbstractArray) = splitview(A, smode) + """ struct NonSplitMode{N} <: AbstractSplitMode From b7b74679483b12ce2866b7c4ab4f184a49540a89 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:15:28 +0200 Subject: [PATCH 20/45] STASH smodes and more with AI --- README.md | 4 +- docs/src/index.md | 24 +- ext/ArraysOfArraysAdaptExt.jl | 6 +- ext/ArraysOfArraysChainRulesCoreExt.jl | 58 ++-- ext/ArraysOfArraysStaticArraysCoreExt.jl | 14 +- src/ArraysOfArrays.jl | 2 +- src/array_of_similar_arrays.jl | 309 +++++++++++---------- src/base_slices.jl | 41 ++- src/functions.jl | 331 ++++++++++++++++------- src/util.jl | 11 +- src/vector_of_arrays.jl | 159 +++++++++-- test/array_of_similar_arrays.jl | 104 +++---- test/base_slices.jl | 8 + test/broadcasting.jl | 6 +- test/static_arrays.jl | 7 +- test/testdefs.jl | 39 +-- test/vector_of_arrays.jl | 71 ++++- 17 files changed, 770 insertions(+), 424 deletions(-) diff --git a/README.md b/README.md index 57b2c1b..efd24e9 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,8 @@ A Julia package for efficient storage and handling of nested arrays. -ArraysOfArrays provides two different types of nested arrays: `SlicedView` and `VectorOfArrays`. -An `SlicedView` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. +ArraysOfArrays provides two different types of nested arrays: `ArrayOfSimilarArrays` and `VectorOfArrays`. +An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. ## Documentation diff --git a/docs/src/index.md b/docs/src/index.md index fb718ed..85ada8f 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,22 +1,26 @@ # ArraysOfArrays.jl -A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: [`SlicedView`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). +A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: [`ArrayOfSimilarArrays`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). This package also defines and exports the following new functions applicable to nested arrays in general: -* [`nestedview`](@ref) and [`flatview`](@ref) switch between a flat and a nested view of the same data. +* [`sliced`](@ref) and [`flatview`](@ref) switch between a flat and a nested (sliced) view of the same data. +* [`partitioned`](@ref) creates a view of a vector as a vector of arrays that may differ in size. +* [`getsplitmode`](@ref), [`splitup`](@ref) and [`fused`](@ref) provide a general API to split arrays into nested (sliced or partitioned) form and to fuse them back into flat form. +* [`stacked`](@ref) and [`unstackmode`](@ref) are similar to `Base.stack`, but can return the original underlying data of sliced arrays without copying it. +* [`vecflattened`](@ref) concatenates the elements of nested arrays into a single vector. * [`innersize`](@ref) returns the size of the elements of an array, provided they all have equal size. * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. -## [SlicedView](@id section_ArrayOfSimilarArrays) +## [ArrayOfSimilarArrays](@id section_ArrayOfSimilarArrays) -An `SlicedView` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays: +An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays: ```julia A_flat = rand(2,3,4,5,6) -A_nested = nestedview(A_flat, 2) +A_nested = sliced(A_flat, 2) ``` creates a view of `A_flat` as an array of arrays: @@ -42,9 +46,9 @@ all(x -> x == 4.2, A_flat[:, :, 2, 4, 3]) The following type aliases are defined: -* `VectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `ArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `VectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `VectorOfSimilarArrays{T,M} = ArrayOfSimilarArrays{T,M,1}` +* `ArrayOfSimilarVectors{T,N} = ArrayOfSimilarArrays{T,1,N}` +* `VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}` For each of the types there is also an abstract type (`AbstractArrayOfSimilarArrays`, etc.). @@ -55,7 +59,7 @@ If a `VectorOfSimilarArrays` is backed by an `ElasticArrays.ElasticArray`, addit ```julia using ElasticArrays -A_nested = nestedview(ElasticArray{Float64}(undef, 2, 3, 0), 2) +A_nested = sliced(ElasticArray{Float64}(undef, 2, 3, 0), 2) for i in 1:4 push!(A_nested, rand(2, 3)) @@ -106,7 +110,7 @@ all(x -> x == 4.2, VA[2]) ### Type aliases The following type aliases are defined: -* `PartsView{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +* `PartsView{T,VT,VI,VD,ET} = VectorOfArrays{T,1,0,VT,VI,VD,ET}` ### Appending data and resizing diff --git a/ext/ArraysOfArraysAdaptExt.jl b/ext/ArraysOfArraysAdaptExt.jl index 6ae9d2d..6bd2e65 100644 --- a/ext/ArraysOfArraysAdaptExt.jl +++ b/ext/ArraysOfArraysAdaptExt.jl @@ -5,13 +5,13 @@ module ArraysOfArraysAdaptExt import Adapt using Adapt: adapt -using ArraysOfArrays: SlicedView, VectorOfArrays +using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays using ArraysOfArrays: no_consistency_checks -function Adapt.adapt_structure(to, A::SlicedView{T,M,N}) where {T,M,N} +function Adapt.adapt_structure(to, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} adapted_data = adapt(to, A.data) - SlicedView{eltype(adapted_data),M,N}(adapted_data) + ArrayOfSimilarArrays{eltype(adapted_data),M,N}(adapted_data) end diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index 3312754..c46ef92 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -4,9 +4,10 @@ module ArraysOfArraysChainRulesCoreExt using ChainRulesCore: ChainRulesCore, NoTangent, AbstractThunk, Thunk, unthunk, @thunk, @non_differentiable -using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitview, fused, flatview, innersize -using ArraysOfArrays: NonSplitMode, AbstractSlicingMode -using ArraysOfArrays: SlicedView +using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitup, fused, stacked, unstackmode, + flatview, innersize +using ArraysOfArrays: NonSplitMode, AbstractSplitMode +using ArraysOfArrays: AbstractArrayOfSimilarArrays, ArrayOfSimilarArrays struct _MappedMaybeThunk{F, T} <: AbstractThunk @@ -19,53 +20,64 @@ mapthunk(::Type{F}, x::T) where {F,T} = _MappedMaybeThunk{Type{F},T}(F, x) @non_differentiable getsplitmode(::Any) +@non_differentiable unstackmode(::Any) @non_differentiable innersize(::Any) @non_differentiable is_memordered_splitmode(::Any) -function ChainRulesCore.rrule(::typeof(splitview), A::AbstractArray, ::NonSplitMode) - return fused(A), _partview_pullback +function ChainRulesCore.rrule(::typeof(splitup), A::AbstractArray, smode::NonSplitMode) + return splitup(A, smode), _nosplit_pullback end -_unpart_partview_pullback(ΔΩ) = NoTangent(), ΔΩ, NoTangent() +_nosplit_pullback(ΔΩ) = NoTangent(), ΔΩ, NoTangent() -function ChainRulesCore.rrule(::typeof(fused), A::AbstractArray) - return fused(A), _unpart_joinedview_pullback +function ChainRulesCore.rrule(::typeof(splitup), A::AbstractArray, smode::AbstractSplitMode) + return splitup(A, smode), _splitview_pullback end -_unpart_joinedview_pullback(ΔΩ) = NoTangent(), ΔΩ +_splitview_pullback(ΔΩ) = NoTangent(), mapthunk(fused, ΔΩ), NoTangent() -function ChainRulesCore.rrule(::typeof(splitview), A::AbstractArray, smode::AbstractSlicingMode) - return splitview(A, smode), _partview_pullback -end -_partview_pullback(ΔΩ) = NoTangent(), mapthunk(fused, ΔΩ), NoTangent() + +ChainRulesCore.rrule(::typeof(fused), A::AbstractArray) = fused(A), _nofuse_pullback +_nofuse_pullback(ΔΩ) = NoTangent(), ΔΩ function ChainRulesCore.rrule(::typeof(fused), A::AbstractArray{<:AbstractArray}) - smode = getsplitmode(A) - return fused(A), Base.Fix2(_joinedview_pullback, smode) + return fused(A), Base.Fix2(_fused_pullback, getsplitmode(A)) +end +_fused_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitup, smode), ΔΩ) + + +ChainRulesCore.rrule(::typeof(stacked), A::AbstractArray) = stacked(A), _nostack_pullback +_nostack_pullback(ΔΩ) = NoTangent(), ΔΩ + +function ChainRulesCore.rrule(::typeof(stacked), A::AbstractArray{<:AbstractArray}) + return stacked(A), Base.Fix2(_stacked_pullback, unstackmode(A)) +end +function ChainRulesCore.rrule(::typeof(stack), A::AbstractArrayOfSimilarArrays) + return stack(A), Base.Fix2(_stacked_pullback, unstackmode(A)) end -_joinedview_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitview, smode), ΔΩ) +_stacked_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitup, smode), ΔΩ) function _aosa_ctor_fromflat_pullback(ΔΩ) - NoTangent(), flatview(convert(SlicedView, unthunk(ΔΩ))) + NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) end -function ChainRulesCore.rrule(::Type{SlicedView{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} - return SlicedView{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback +function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} + return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback end _aosa_ctor_fromnested_pullback(ΔΩ) = NoTangent(), ΔΩ -function ChainRulesCore.rrule(::Type{SlicedView{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} - return SlicedView{T,M,N}(A), _aosa_ctor_fromnested_pullback +function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} + return ArrayOfSimilarArrays{T,M,N}(A), _aosa_ctor_fromnested_pullback end -function ChainRulesCore.rrule(::typeof(flatview), A::SlicedView{T,M,N}) where {T,M,N} +function ChainRulesCore.rrule(::typeof(flatview), A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} function flatview_pullback(ΔΩ) data = unthunk(ΔΩ) - NoTangent(), SlicedView{eltype(data),M,N}(data) + NoTangent(), ArrayOfSimilarArrays{eltype(data),M,N}(data) end return flatview(A), flatview_pullback diff --git a/ext/ArraysOfArraysStaticArraysCoreExt.jl b/ext/ArraysOfArraysStaticArraysCoreExt.jl index c370949..62094a3 100644 --- a/ext/ArraysOfArraysStaticArraysCoreExt.jl +++ b/ext/ArraysOfArraysStaticArraysCoreExt.jl @@ -6,21 +6,27 @@ import StaticArraysCore using StaticArraysCore: StaticArray, SVector import ArraysOfArrays -using ArraysOfArrays: nestedview +using ArraysOfArrays: sliced @inline ArraysOfArrays.flatview(A::AbstractArray{SA,N}) where {S,T,M,N,SA<:StaticArray{S,T,M}} = reshape(reinterpret(T, A), size(SA)..., size(A)...) -@inline function ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} +@inline function ArraysOfArrays.sliced(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} size_A = size(A) size_A[1] == S || throw(DimensionMismatch("Length $S of static vector type does not match first dimension of array of size $size_A")) reshape(reinterpret(SA, A), ArraysOfArrays._tail(size_A)...) end -@inline ArraysOfArrays.nestedview(A::AbstractArray{T}, ::Type{SVector{S}}) where {T,S} = - nestedview(A, SVector{S,T}) +@inline ArraysOfArrays.sliced(A::AbstractArray{T}, ::Type{SVector{S}}) where {T,S} = + sliced(A, SVector{S,T}) + + +# Deprecated: + +Base.@deprecate ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} ArraysOfArrays.sliced(A, SA) false +Base.@deprecate ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S}}) where {T,S} ArraysOfArrays.sliced(A, SA) false end # module ArraysOfArraysStaticArraysCoreExt diff --git a/src/ArraysOfArrays.jl b/src/ArraysOfArrays.jl index 06973dd..2a0788d 100644 --- a/src/ArraysOfArrays.jl +++ b/src/ArraysOfArrays.jl @@ -8,7 +8,7 @@ __precompile__(true) Efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: -[`SlicedView`](@ref section_ArrayOfSimilarArrays) and +[`ArrayOfSimilarArrays`](@ref section_ArrayOfSimilarArrays) and [`VectorOfArrays`](@ref section_VectorOfArrays). """ module ArraysOfArrays diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 5fd92c4..2ff7a0f 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -33,215 +33,227 @@ export AbstractVectorOfSimilarVectors """ - SlicedView{T,M,N,P} <: AbstractArrayOfSimilarArrays{T,M,N} + struct SplitSlices{M,N} <: AbstractSlicingMode{M,N} -Represents a view of an array of dimension `M + N` as an array of -dimension M with elements that are arrays with dimension N. All element arrays +The split mode of [`ArrayOfSimilarArrays`](@ref): memory-ordered slicing +with `M` inner and `N` outer dimensions. + +Constructor: + +``` +SplitSlices{M,N}() +``` + +See also [`AbstractSlicingMode`](@ref). +""" +struct SplitSlices{M,N} <: AbstractSlicingMode{M,N} end +export SplitSlices + +is_memordered_splitmode(::SplitSlices) = true + +getinnerdims(obj::Tuple, ::SplitSlices{M,N}) where {M,N} = front_tuple(obj, Val(M)) +getouterdims(obj::Tuple, ::SplitSlices{M,N}) where {M,N} = back_tuple(obj, Val(N)) + +@inline splitup(A::AbstractArray{T}, ::SplitSlices{M,N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(A) + + +""" + ArrayOfSimilarArrays{T,M,N,P,ET} <: AbstractArrayOfSimilarArrays{T,M,N,ET} + +Represents a view of an array of dimension `M + N` as an `N`-dimensional +array with elements that are `M`-dimensional arrays. All element arrays implicitly have equal size/axes. Constructors: - SlicedView{T,M,N}(flat_data::AbstractArray) - SlicedView{T,M}(flat_data::AbstractArray) + ArrayOfSimilarArrays{T,M,N}(data::AbstractArray) + ArrayOfSimilarArrays{T,M}(data::AbstractArray) The following type aliases are defined: -* `VectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `ArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `VectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `VectorOfSimilarArrays{T,M} = ArrayOfSimilarArrays{T,M,1}` +* `ArrayOfSimilarVectors{T,N} = ArrayOfSimilarArrays{T,1,N}` +* `VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}` `VectorOfSimilarArrays` supports `push!()`, etc., provided the underlying array supports resizing of it's last dimension (e.g. an `ElasticArray`). -The nested array can also be created using the function [`nestedview`](@ref) +The nested array can also be created using the function [`sliced`](@ref) and the wrapped flat array can be accessed using [`flatview`](@ref) afterwards: ```julia A_flat = rand(2,3,4,5,6) -A_nested = nestedview(A_flat, 2) +A_nested = sliced(A_flat, Val(2)) A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T flatview(A_nested) === A_flat ``` """ -struct SlicedView{ +struct ArrayOfSimilarArrays{ T, M, N, P<:AbstractArray{T}, ET<:AbstractArray{T,M} } <: AbstractArrayOfSimilarArrays{T,M,N,ET} data::P - function SlicedView{T,M,N}(flat_data::AbstractArray{U}) where {T,M,N,U} - require_ndims(flat_data, _add_vals(Val{M}(), Val{N}())) - conv_parent = _convert_elype(T, flat_data) - P = typeof(conv_parent) - ET = Base.promote_op(view, P, _nColons(Val{M}())..., _nInts(Val{N}())...) - new{T,M,N,P,ET}(conv_parent) + function ArrayOfSimilarArrays{T,M,N}(data::AbstractArray{T,L}) where {T,M,N,L} + _require_ndims(Val(L), _add_vals(Val(M), Val(N))) + P = typeof(data) + ET = Base.promote_op(view, P, _nColons(Val(M))..., _nInts(Val(N))...) + new{T,M,N,P,ET}(data) end end -function SlicedView{T,M}(flat_data::AbstractArray{U}) where {T,M,U} - _, size_outer = split_tuple(size(flat_data), Val{M}()) - N = length(size_outer) - SlicedView{T,M,N}(flat_data) -end +export ArrayOfSimilarArrays -export SlicedView -function SlicedView{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} - B = SlicedView{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...)) - copyto!(B, A) +function ArrayOfSimilarArrays{T,M,N}(orig_data::AbstractArray{U,L}) where {T,M,N,U,L} + conv_data = _convert_elype(T, orig_data)::AbstractArray{T,L} + return ArrayOfSimilarArrays{T,M,N}(conv_data) end -SlicedView{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - SlicedView{T,M,N}(A) +function ArrayOfSimilarArrays{T,M}(data::AbstractArray{U,L}) where {T,M,U,L} + N = _val_value(_subtract_vals(Val(L), Val(M))) + ArrayOfSimilarArrays{T,M,N}(data) +end + +Base.convert(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) +Base.convert(::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) +Base.convert(::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) + + +@inline Base.:(==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} = (stacked(A) == stacked(B)) +@inline Base.isequal(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} = isequal(stacked(A), stacked(B)) +@inline Base.isapprox(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}; kwargs...) where {T,M,N,U} = isapprox(stacked(A), stacked(B); kwargs...) -SlicedView(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = - SlicedView{T,M,N}(A) +@inline getsplitmode(A::ArrayOfSimilarArrays) = unstackmode(A) +@inline unstackmode(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = SplitSlices{M,N}() +@inline stacked(A::ArrayOfSimilarArrays) = A.data +@inline fused(A::ArrayOfSimilarArrays) = stacked(A) -Base.convert(R::Type{SlicedView{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} = R(flat_data) -Base.convert(R::Type{SlicedView{T,M}}, flat_data::AbstractArray{U}) where {T,M,U} = R(flat_data) +@inline vecflattened(A::ArrayOfSimilarArrays) = vec(fused(A)) -Base.convert(R::Type{SlicedView{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{SlicedView{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{SlicedView}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) +@inline Base.parent(A::ArrayOfSimilarArrays) = fused(A) -fused(A::SlicedView{T,M,N}) where {T,M,N} = A.data -Base.stack(A::SlicedView) = fused(A) +# `stack` must return an independent array, unlike `stacked`. Julia does not +# dispatch on keyword arguments, so dispatch on the value of `dims` instead: +Base.stack(A::ArrayOfSimilarArrays; dims::Union{Integer,Colon} = :) = _stack_impl(A, dims) -function Base.Array(A::SlicedView{T,M,N,P,ET}) where {T,M,N,P,ET} +# Fast path for the default layout: a single bulk copy instead of the +# element-by-element copy that generic `stack` would do: +_stack_impl(A::ArrayOfSimilarArrays, ::Colon) = copy(fused(A)) +_stack_impl(A::ArrayOfSimilarArrays, dims::Integer) = stack(collect(A); dims) + +function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} new_ET = Base.promote_op(similar, ET) return Array{new_ET,N}(A) end -function getslicemap(::SlicedView{T,M,N}) where {T,M,N} + +function getslicemap(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) end -#!!!!! -@inline function innersize(A::SlicedView{T,M,N}) where {T,M,N} + +@inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} front_tuple(size(A.data), Val{M}()) end -import Base.== -(==)(A::SlicedView{T,M,N}, B::SlicedView{T,M,N}) where {T,M,N} = - (A.data == B.data) - +Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] -Base.size(A::SlicedView{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] - - -Base.@propagate_inbounds Base.getindex(A::SlicedView{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} = +Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} = view(A.data, _ncolons(Val{M}())..., idxs...) -Base.@propagate_inbounds Base.setindex!(A::SlicedView{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = +Base.@propagate_inbounds Base.setindex!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = setindex!(A.data, x, _ncolons(Val{M}())..., idxs...) -Base.@propagate_inbounds function Base.unsafe_view(A::SlicedView{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} +Base.@propagate_inbounds function Base.unsafe_view(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} dataview = view(A.data, _ncolons(Val{M}())..., idxs...) L = length(size(dataview)) N_view = L - M - SlicedView{T,M,N_view}(dataview) + ArrayOfSimilarArrays{T,M,N_view}(dataview) end -@inline function Base.resize!(A::SlicedView{T,M,N}, dims::Vararg{Integer,N}) where {T,M,N} +@inline function Base.resize!(A::ArrayOfSimilarArrays{T,M,N}, dims::Vararg{Integer,N}) where {T,M,N} resize!(A.data, innersize(A)..., dims...) A end -function Base.similar(A::SlicedView{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U} +function Base.similar(A::ArrayOfSimilarArrays{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U} data = A.data - size_inner, size_outer = split_tuple(size(data), Val{M}()) + size_inner = front_tuple(size(data), Val{M}()) # ToDo: Don't use similar if data is an ElasticArray? - SlicedView{T,M,N}(similar(data, U, size_inner..., dims...)) + ArrayOfSimilarArrays{T,M,N}(similar(data, U, size_inner..., dims...)) end -function Base.deepcopy(A::SlicedView{T,M,N}) where {T,M,N} - SlicedView{T,M,N}(deepcopy(A.data)) +function Base.deepcopy(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} + ArrayOfSimilarArrays{T,M,N}(deepcopy(A.data)) end -function Base.copyto!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} +function Base.copyto!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} copyto!(dest.data, src.data) dest end -function Base.append!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} +function Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't append, shape of element arrays of source and dest are not equal")) append!(dest.data, src.data) dest end -Base.append!(dest::SlicedView{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - append!(dest, SlicedView(src)) +Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = + append!(dest, ArrayOfSimilarArrays(src)) -function Base.prepend!(dest::SlicedView{T,M,N}, src::SlicedView{U,M,N}) where {T,M,N,U} +function Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't prepend, shape of element arrays of source and dest are not equal")) prepend!(dest.data, src.data) dest end -Base.prepend!(dest::SlicedView{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - prepend!(dest, SlicedView(src)) +Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = + prepend!(dest, ArrayOfSimilarArrays(src)) -function innermap(f::Base.Callable, A::SlicedView{T,M,N}) where {T,M,N} +function innermap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} new_data = map(f, A.data) U = eltype(new_data) - SlicedView{U,M,N}(new_data) + ArrayOfSimilarArrays{U,M,N}(new_data) end -function deepmap(f::Base.Callable, A::SlicedView{T,M,N}) where {T,M,N} +function deepmap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} new_data = deepmap(f, A.data) U = eltype(new_data) - SlicedView{U,M,N}(new_data) + ArrayOfSimilarArrays{U,M,N}(new_data) end -Base.map(::typeof(identity), A::SlicedView) = A -Base.Broadcast.broadcasted(::typeof(identity), A::SlicedView) = A - - -Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) = - Val{!(Base.index_dimsum(idxs_outer...) isa Tuple{}) && !(Base.index_dimsum(idxs_inner...) isa Tuple{})}() - -Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs...))}() +Base.map(::typeof(identity), A::ArrayOfSimilarArrays) = A +Base.Broadcast.broadcasted(::typeof(identity), A::ArrayOfSimilarArrays) = A const VectorOfSimilarArrays{ T, M, P<:AbstractArray{T}, - ET<:AbstractArray{T} -} = SlicedView{T,M,1,P,ET} + ET<:AbstractArray{T,M} +} = ArrayOfSimilarArrays{T,M,1,P,ET} export VectorOfSimilarArrays -VectorOfSimilarArrays{T}(flat_data::AbstractArray{U}) where {T,U} = - SlicedView{T,length(Base.front(size(flat_data))),1}(flat_data) - -VectorOfSimilarArrays(flat_data::AbstractArray{T}) where {T} = - SlicedView{T,length(Base.front(size(flat_data))),1}(flat_data) +VectorOfSimilarArrays{T}(data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,length(Base.front(size(data))),1}(data) +VectorOfSimilarArrays(data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,length(Base.front(size(data))),1}(data) -VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = - VectorOfSimilarArrays{T,M}(A) - -VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = - VectorOfSimilarArrays{T,M}(A) - - -Base.convert(R::Type{VectorOfSimilarArrays{T}}, flat_data::AbstractArray{U}) where {T,U} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarArrays}, flat_data::AbstractArray{T}) where {T} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = R(A) -Base.convert(R::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = R(A) +Base.convert(::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = ArrayOfSimilarArrays{T,M,1}(stacked(A)) +Base.convert(::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = ArrayOfSimilarArrays{T,M,1}(stacked(A)) @inline Base.IndexStyle(::Type{<:VectorOfSimilarArrays}) = IndexLinear() @@ -267,7 +279,7 @@ function Base.pushfirst!(V::VectorOfSimilarArrays{T,M}, x::AbstractArray{U,M}) w end # Will need equivalent of resize! that resizes in front of data instead of in back: -# popfirst!(V::SlicedView) = ... +# popfirst!(V::ArrayOfSimilarArrays) = ... function _empty_data_size(A::VectorOfSimilarArrays{T,M}) where {T,M} @@ -279,7 +291,7 @@ end function Base.empty(A::VectorOfSimilarArrays{T,M}, ::Type{<:AbstractArray{U}}) where {T,M,U} new_data_size = _empty_data_size(A) # ToDo: Don't use similar if data is an ElasticArray? - VectorOfSimilarArrays{T,M}(similar(A.data, U, new_data_size...)) + ArrayOfSimilarArrays{U,M,1}(similar(A.data, U, new_data_size...)) end function Base.empty!(A::VectorOfSimilarArrays{T,M}) where {T,M} @@ -293,53 +305,30 @@ const ArrayOfSimilarVectors{ T, N, P<:AbstractArray{T}, ET<:AbstractVector{T} -} = SlicedView{T,1,N,P,ET} +} = ArrayOfSimilarArrays{T,1,N,P,ET} export ArrayOfSimilarVectors -ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U}) where {T,U} = - SlicedView{T,1,length(Base.front(size(flat_data)))}(flat_data) - -ArrayOfSimilarVectors(flat_data::AbstractArray{T}) where {T} = - SlicedView{T,1,length(Base.front(size(flat_data)))}(flat_data) +ArrayOfSimilarVectors{T}(data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,1,length(Base.front(size(data)))}(data) +ArrayOfSimilarVectors(data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,1,length(Base.front(size(data)))}(data) -ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = - ArrayOfSimilarVectors{T,N}(A) - -ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = - ArrayOfSimilarVectors{T,N}(A) - - -Base.convert(R::Type{ArrayOfSimilarVectors{T}}, flat_data::AbstractArray{U}) where {T,U} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarVectors}, flat_data::AbstractArray{T}) where {T} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(A) +Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(stacked(A)) +Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(stacked(A)) const VectorOfSimilarVectors{ T, P<:AbstractArray{T,2}, ET<:AbstractVector{T} -} = SlicedView{T,1,1,P,ET} +} = ArrayOfSimilarArrays{T,1,1,P,ET} export VectorOfSimilarVectors -VectorOfSimilarVectors{T}(flat_data::AbstractArray{U,2}) where {T,U} = - SlicedView{T,1,1}(flat_data) - -VectorOfSimilarVectors(flat_data::AbstractArray{T,2}) where {T} = - VectorOfSimilarVectors{T}(flat_data) +VectorOfSimilarVectors{T}(data::AbstractArray{U,2}) where {T,U} = ArrayOfSimilarArrays{T,1,1}(data) +VectorOfSimilarVectors(data::AbstractArray{T,2}) where {T} = VectorOfSimilarVectors{T}(data) -VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} = - SlicedView{T,1}(A) - -VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} = - VectorOfSimilarVectors{T}(A) - -Base.convert(R::Type{VectorOfSimilarVectors{T}}, flat_data::AbstractArray{U,2}) where {T,U} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarVectors}, flat_data::AbstractArray{T,2}) where {T} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVector{U}}) where {T,U} = R(A) -Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(A) +Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVector{U}}) where {T,U} = R(stacked(A)) +Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(stacked(A)) Base.sum(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = @@ -361,43 +350,47 @@ Statistics.cor(X::AbstractVectorOfSimilarVectors) = cor(flatview(X); dims = 2) -""" - nestedview(A::AbstractArray{T,M+N}, M::Integer) - nestedview(A::AbstractArray{T,2}) -AbstractArray{<:AbstractArray{T,M},N} +# Deprecations: -View array `A` in as an `N`-dimensional array of `M`-dimensional arrays by -wrapping it into an [`SlicedView`](@ref). +@deprecate ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false +@deprecate ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false +@deprecate ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false -It's also possible to use a `StaticVector` of length `S` as the type of the -inner arrays via +@deprecate Base.convert(::Type{ArrayOfSimilarArrays{T,M,N}}, data::AbstractArray{U,L}) where {T,M,N,U<:Number,L} ArrayOfSimilarArrays{T,M,N}(data) false +@deprecate Base.convert(::Type{ArrayOfSimilarArrays{T,M}}, data::AbstractArray{U,L}) where {T,M,U<:Number,L} ArrayOfSimilarArrays{T,M}(data) false - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S}}) - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S,T}}) -""" -function nestedview end -export nestedview -@inline nestedview(A::AbstractArray{T,L}, M::Integer) where {T,L} = - SlicedView{T,M}(A) +@deprecate VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} ArrayOfSimilarArrays{T,M,1}(stacked(A)) false +@deprecate VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} ArrayOfSimilarArrays{T,M,1}(stacked(A)) false -@inline nestedview(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = - SlicedView{T,M}(A) +@deprecate Base.convert(::Type{VectorOfSimilarArrays{T}}, data::AbstractArray{U}) where {T,U<:Number} VectorOfSimilarArrays{T}(data) false +@deprecate Base.convert(::Type{VectorOfSimilarArrays}, data::AbstractArray{T}) where {T<:Number} VectorOfSimilarArrays(data) false -@inline nestedview(A::AbstractArray{T,2}) where {T} = - VectorOfSimilarVectors{T}(A) +@deprecate ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} ArrayOfSimilarArrays{T,1,N}(stacked(A)) false +@deprecate ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} ArrayOfSimilarArrays{T,1,N}(stacked(A)) false -# Deprecated: -const ArrayOfSimilarArrays{T,M,N} = SlicedView{T,M,N} -export ArrayOfSimilarArrays +@deprecate Base.convert(::Type{ArrayOfSimilarVectors{T}}, data::AbstractArray{U}) where {T,U<:Number} ArrayOfSimilarVectors{T}(data) false +@deprecate Base.convert(::Type{ArrayOfSimilarVectors}, data::AbstractArray{T}) where {T<:Number} ArrayOfSimilarVectors(data) false -const VectorOfSimilarArrays{T,M} = VoASlicedView{T,M} -export VectorOfSimilarArrays -const ArrayOfSimilarVectors{T,N} = AoVSlicedView{T,N} -export ArrayOfSimilarVectors +@deprecate VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} ArrayOfSimilarArrays{T,1,1}(stacked(A)) false +@deprecate VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} ArrayOfSimilarArrays{T,1,1}(stacked(A)) false -const VectorOfSimilarVectors{T} = VoVSlicedView{T} -export VectorOfSimilarVectors +@deprecate Base.convert(::Type{VectorOfSimilarVectors{T}}, data::AbstractArray{U,2}) where {T,U<:Number} VectorOfSimilarVectors{T}(data) false +@deprecate Base.convert(::Type{VectorOfSimilarVectors}, data::AbstractArray{T,2}) where {T<:Number} VectorOfSimilarVectors(data) false + + +""" + nestedview(A::AbstractArray{T,M+N}, M::Integer) + nestedview(A::AbstractArray{T,2}) + +Deprecated, use [`sliced`](@ref) instead. +""" +function nestedview end +export nestedview + +@deprecate nestedview(A::AbstractArray, M::Integer) sliced(A, Val(M)) +@deprecate nestedview(A::AbstractArray, ::Val{M}) where {M} sliced(A, Val(M)) +@deprecate nestedview(A::AbstractArray{T,2}) where {T} sliced(A) diff --git a/src/base_slices.jl b/src/base_slices.jl index 6fd3aa6..34e201a 100644 --- a/src/base_slices.jl +++ b/src/base_slices.jl @@ -1,17 +1,18 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). """ - struct BaseSlicing{N,TPL<:Tuple{Vararg{Int,N}}} <: AbstractSlicingMode + struct BaseSlicing{M,N,TPL<:Tuple{Vararg{Union{Colon,Int}}}} <: AbstractSlicingMode{M,N} -The split mode of `Slices`. +The split mode of `Base.Slices` (as returned by `eachslice`, `eachcol` and +`eachrow`). Constructor: ``` -BaseSlicing(slicemap::Tuple{Vararg{Union{Colon,Int}}}}) +BaseSlicing{M,N,TPL}(slicemap::TPL) ``` -`slicemap` equals the `slicemap` property of `Base.Slice` objects. +`slicemap` equals the `slicemap` property of `Base.Slices` objects. See also [`AbstractSlicingMode`](@ref). """ @@ -49,27 +50,43 @@ end end end - outdimidxs_expr = Expr(:tuple, outdimidxs...) - idxorder_expr = Expr(:tuple, [:(slicemap[$i]) for i in outdimidxs]...) - result_expr = Expr(:tuple, [:(obj[outdimidxs[idxorder[$i]]]) for i in eachindex(outdimidxs)]...) + # With less than two outer dimensions the result is fully determined by + # SliceMapT already: + if length(outdimidxs) < 2 + return Expr(:tuple, [:(obj[$i]) for i in outdimidxs]...) + end + + # vals[j] is the entry of obj and pos[j] the outer dimension number for + # the j-th outer dimension of the parent array. The result must be + # ordered by outer dimension number, so it is vals in inverse-perm order + # of pos. pos is only known at runtime, its values are not encoded in + # SliceMapT. + vals_expr = Expr(:tuple, [:(obj[$i]) for i in outdimidxs]...) + pos_expr = Expr(:tuple, [:(slicemap[$i]) for i in outdimidxs]...) quote slicemap = smode.slicemap - outdimidxs = $outdimidxs_expr - idxorder = $idxorder_expr - return $result_expr + vals = $vals_expr + pos = $pos_expr + return _invpermuted(vals, pos) end end +@inline function _invpermuted(vals::NTuple{N,Any}, pos::NTuple{N,Int}) where N + ntuple(k -> vals[findfirst(==(k), pos)::Int], Val(N)) +end + + +@inline getslicemap(A::Slices) = A.slicemap @inline function getsplitmode(A::Slices) M = ndims(eltype(A)) N = ndims(A) - slicemap = A.slicemap + slicemap = getslicemap(A) BaseSlicing{M,N,typeof(slicemap)}(slicemap) end -function splitview(A::AbstractArray, smode::BaseSlicing) +function splitup(A::AbstractArray, smode::BaseSlicing) slicemap = smode.slicemap axs = getouterdims(axes(A), smode) return Slices(A, slicemap, axs) diff --git a/src/functions.jl b/src/functions.jl index fccfac5..694a4b0 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -8,29 +8,44 @@ Abstract supertype for array split modes. Use [`getsplitmode`](@ref) to get the split mode of an array. -Use [`splitview`](@ref) or call `smode::AbstractSplitMode` as a function to +Use [`splitup`](@ref) or call `smode::AbstractSplitMode` as a function to split an array: ```julia -splitview(A, smode) === smode(A) +splitup(A, smode) === smode(A) ``` -See also [`splitview`](@ref) and [`fused`](@ref). +See also [`splitup`](@ref) and [`fused`](@ref). # Implementation Subtypes of `AbstractSplitMode` should specialize -* `splitview(A, smode::SomeSplitMode)` for arrays `A` +* `splitup(A, smode::SomeSplitMode)` for arrays `A` * `is_memordered_splitmode(smode::SomeSplitMode)` -`(smode::SomeSplitMode)(A)` calls `splitview(A, smode)` by default and should +`(smode::SomeSplitMode)(A)` calls `splitup(A, smode)` by default and should not be specialized. """ abstract type AbstractSplitMode <: Function end export AbstractSplitMode -(smode::AbstractSplitMode)(A::AbstractArray) = splitview(A, smode) +(smode::AbstractSplitMode)(A::AbstractArray) = splitup(A, smode) + + +""" + is_memordered_splitmode(smode::AbstractSplitMode)::Bool + +Check if `smode` splits in memory-order. + +If true, inner arrays are stored contiguously in memory in Julia-native +dimension order, and the same is true for the outer dimensions (no dimention +reordering). + +If true, `flatview` and `fused` are equivalent. +""" +function is_memordered_splitmode end +export is_memordered_splitmode """ @@ -43,16 +58,18 @@ Constructor: `NonSplitMode{N}()` struct NonSplitMode{N} <: AbstractSplitMode end export NonSplitMode +is_memordered_splitmode(::NonSplitMode) = true + """ struct UnknownSplitMode{AT} <: AbstractSplitMode -Split mode of generic split objects of type{T} that have been split in an +Split mode of generic split objects of type `AT` that have been split in an unknown way, e.g. nested arrays of type `Array{<:Array}`. Since the split parts may be (typically are) non-contiguous in memory, this -split mode not allow for `unsplitview` or `flatview`. It is also not -inferrable of the split object should be interpeted as a sliced array, +split mode does not allow for `fused` or `flatview`. It is also not +inferrable if the split object should be interpreted as a sliced array, a ragged array, or something else. Constructor: `UnknownSplitMode{T}()` @@ -60,18 +77,7 @@ Constructor: `UnknownSplitMode{T}()` struct UnknownSplitMode{AT} <: AbstractSplitMode end export UnknownSplitMode - -""" - abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode - -Abstract supertype for array slicing modes with `M` inner dimensions and `N` -outer dimensions. - -Use `getsplitmode` to get the split mode of an split array. -""" -abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode end -export AbstractSlicingMode - +is_memordered_splitmode(::UnknownSplitMode) = false """ @@ -80,7 +86,7 @@ export AbstractSlicingMode Get the split mode of `A`. -`splitview(fused(A), getsplitmode(A))` must equal `A`, and should have +`splitup(fused(A), getsplitmode(A))` must equal `A`, and should have the same type as `A` if at all possible, except if `getsplitmode(A)` is an `UnknownSplitMode`. @@ -97,69 +103,21 @@ export getsplitmode """ - is_memordered_splitmode(smode::AbstractSplitMode)::Bool - -Check if `smode` splits in memory-order. - -If true, inner arrays are stored contiguously in memory in Julia-native -dimension order, and the same is true for the outer dimensions (no dimention -reordering). - -If true, `flatview` and `fused` are equivalent. -""" -function is_memordered_splitmode end -export is_memordered_splitmode - -is_memordered_splitmode(::NonSplitMode) = true -is_memordered_splitmode(::UnknownSplitMode) = false - - -""" - ArraysOfArrays.getinnerdims(tpl::Tuple, smode::AbstractSlicing) - -Get the entries of `tpl` corresponding to the inner dimensions of slicing -mode `smode`, in the order specified by `smode`. -""" -function getinnerdims end - -@inline getinnerdims(::Tuple, ::NonSplitMode) = () - -function getinnerdims(::Tuple, ::UnknownSplitMode) - throw(ArgumentError("getinnerdims cannot be used with UnknownSplitMode")) -end - - -""" - ArraysOfArrays.getouterdims(tpl::Tuple, smode::AbstractSlicing) - -Get the entries of `tpl` corresponding to the outer dimensions of slicing -mode `smode`, in the order specified by `smode`. -""" -function getouterdims end - -@inline getouterdims(x::Tuple, ::NonSplitMode) = x - -function getouterdims(::Tuple, ::UnknownSplitMode) - throw(ArgumentError("getouterdims cannot be used with UnknownSplitMode")) -end - - -""" - splitview(A::AbstractArray, smode::AbstractSplitMode) + splitup(A::AbstractArray, smode::AbstractSplitMode) View array `A` in split form, as an array of arrays. -`splitview` should be a zero-copy O(1) operation, if at all possible. +`splitup` should be a zero-copy O(1) operation, if at all possible. See also [`fused`](@ref) and [`getsplitmode`](@ref). """ -function splitview end -export splitview +function splitup end +export splitup -@inline splitview(obj::Any, ::NonSplitMode) = obj +@inline splitup(obj::Any, ::NonSplitMode) = obj -function splitview(::Any, ::UnknownSplitMode) - throw(ArgumentError("splitview cannot be used with UnknownSplitMode")) +function splitup(::Any, ::UnknownSplitMode) + throw(ArgumentError("splitup cannot be used with UnknownSplitMode")) end @@ -169,7 +127,7 @@ end View array `A` in unsplit form. -`splitview(fused(A), getsplitmode(A))` must equal `A`, and should have +`splitup(fused(A), getsplitmode(A))` must equal `A`, and should have the same type as `A` if at all possible, except if `getsplitmode(A)` is an `UnknownSplitMode`. @@ -184,15 +142,15 @@ equivalent to [`flatview(A)`](@ref). function fused end export fused -@inline fused(obj) = _joinedview_impl(obj, getsplitmode(obj)) +@inline fused(obj) = _fused_impl(obj, getsplitmode(obj)) @inline fused(A::AbstractArray) = A -@inline fused(A::AbstractArray{<:AbstractArray}) = _joinedview_impl(A, getsplitmode(A)) +@inline fused(A::AbstractArray{<:AbstractArray}) = _fused_impl(A, getsplitmode(A)) -@inline _joinedview_impl(obj, ::NonSplitMode) = obj +@inline _fused_impl(obj, ::NonSplitMode) = obj -function _joinedview_impl(@nospecialize(obj), ::UnknownSplitMode) +function _fused_impl(@nospecialize(obj), ::UnknownSplitMode) throw(ArgumentError("fused not implemented for objects of type $(nameof(typeof(obj))) with unknown split mode")) end @@ -206,14 +164,14 @@ the flattened form will depend on the type of `A`. If the `A` is not a nested array, the return value is `A` itself. Only specific types of nested arrays are supported. +`flatview` is a zero-copy O(1) operation. + If `is_memordered_splitmode(getsplitmode(A))` is true, `flatview(A)` is equivalent to [`fused(A)`](@ref). -The result of `flatview(A)` will equal either `stack(A)` -(resp. [`stacked(A)`](@ref)) or `reduce(vcat, A)`, depending on the type of -`A` (sliced-array-like or ragged-array-like). - -`flatview` should be a zero-copy O(1) operation, if at all possible. +For sliced arrays the result of `flatview(A)` will equal [`stacked(A)`](@ref). +For partitioned vectors it will equal [`vecflattened(A)`](@ref), provided +that the parts cover the underlying data completely. """ function flatview end export flatview @@ -223,18 +181,105 @@ function flatview(A::AbstractArray{<:AbstractArray}) throw(ArgumentError("flatview not implemented for nested arrays of type $(nameof(typeof(A)))")) end -function flatview(A::AbstractSlices) +function flatview(A::AbstractSlices{<:AbstractArray}) smode = getsplitmode(A) if is_memordered_splitmode(smode) return fused(A) else - throw(ArgumentError("flatview required memory-ordered split/slicing, but array has split mode $smode")) + throw(ArgumentError("flatview requires memory-ordered split/slicing, but array has split mode $smode")) end end + + +""" + abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode + +Abstract supertype for array slicing modes with `M` inner dimensions and `N` +outer dimensions. + +Use `getsplitmode` to get the split mode of an split array. +""" +abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode end +export AbstractSlicingMode + + +""" + getslicemap(A::AbstractSlices) + +Return the slicemap of `A` with respect to `B = fused(A)`: a tuple with one +entry per dimension of `B`, `Colon()` for sliced (inner) dimensions and `k` +for dimensions indexed by dimension `k` of `A`, so that + +```julia +A[i...] == view(B, map(s -> s isa Colon ? (:) : i[s], getslicemap(A))...) +``` + +E.g. `A = eachslice(B, dims = (3,1,5))` of a five-dimensional `B` has the +slicemap `(2, :, 1, :, 3)`, since `A[i1, i2, i3] == view(B, i2, :, i1, :, i3)`. + +Equals the `slicemap` field of `Base.Slices` objects. +""" +function getslicemap end +export getslicemap + + +""" + ArraysOfArrays.getinnerdims(tpl::Tuple, smode::AbstractSplitMode) + +Get the entries of `tpl` corresponding to the inner dimensions of split +mode `smode`, in the order specified by `smode`. +""" +function getinnerdims end + +@inline getinnerdims(::Tuple, ::NonSplitMode) = () + +function getinnerdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getinnerdims cannot be used with UnknownSplitMode")) +end + + """ - stacked(A::AbstractArray) - stacked(A::AbstractArray{<:AbstractArray}) + ArraysOfArrays.getouterdims(tpl::Tuple, smode::AbstractSplitMode) + +Get the entries of `tpl` corresponding to the outer dimensions of split +mode `smode`, in the order specified by `smode`. +""" +function getouterdims end + +@inline getouterdims(x::Tuple, ::NonSplitMode) = x + +function getouterdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getouterdims cannot be used with UnknownSplitMode")) +end + + +""" + unstackmode(A::AbstractArray) + unstackmode(A::AbstractArray{<:AbstractArray}) + +Get the split mode required to restore `stacked(A)` so that +`splitup(stacked(A), unstackmode(A)) == A`. + +The result of `splitup(stacked(A), unstackmode(A))` may have a different +type and underlying memory layout than `A`. +""" +function unstackmode end +export unstackmode + +@inline unstackmode(::T) where T = UnknownSplitMode{T}() + +@inline unstackmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() + +function unstackmode(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} + innersize(A) # Ensure element arrays have equal size + return SplitSlices{M,N}() +end + + +""" + stacked(A::AbstractArray{T,N})::AbstractArray{T,N} + stacked(A::AbstractArray{<:AbstractArray{T,M},N})::AbstractArray{T,M+N} Join stacked arrays of a nested array into a single array along one or more new dimensions, return non-nested arrays unchanged. @@ -250,23 +295,32 @@ export stacked _stacked_impl(A::AbstractArray{<:AbstractArray}, ::AbstractSplitMode) = stack(A) -function _stacked_impl(A::AbstractSlices, smode::AbstractSlicingMode) +function _stacked_impl(A::AbstractSlices{<:AbstractArray}, smode::AbstractSlicingMode) A_joined = fused(A) is_memordered_splitmode(smode) ? A_joined : _stacked_permutedims(A_joined, smode) end function _stacked_permutedims(A_joined::AbstractArray{T,N}, smode::AbstractSlicingMode) where {T,N} dimnumbers = _oneto_tpl(Val(N)) - dimorder = (getinnerdims(dimnumbers, smode)..., getouterdims(_dimnumbers, smode)...) - return permutedims(A_joined, dimorder)::typeof(A_joined) + dimorder = (getinnerdims(dimnumbers, smode)..., getouterdims(dimnumbers, smode)...) + return permutedims(A_joined, dimorder) end -# ToDo: Add function `unstacked(A::AbstractArray, smode::AbstractSlicingMode)`? +""" + sliced(A::AbstractArray{T,2}) + sliced(A::AbstractArray{T,M+N}, Val(M)) + sliced(A::AbstractArray{T,M+N}, M::Integer) + +Return a sliced view of `A`, using the columns or the the first `M` +dimensions as inner dimensions. +""" +function sliced end +export sliced -# ToDo: Add function `stacked_splitmode(smode::AbstractSlicingMode)`, and -# `stacked_splitmode(A::AbstractArray{<:AbstractArray}) = stacked_splitmode(getsplitmode(A))`, -# so that `splitview(stacked(A), stacked_splitmode(A) == A`? +@inline sliced(A::AbstractArray, M::Integer) = sliced(A, Val(M)) +@inline sliced(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = splitup(A, SplitSlices{M,L-M}()) +@inline sliced(A::AbstractArray{T,2}) where {T} = sliced(A, Val(1)) """ @@ -300,7 +354,7 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} return s end -@inline innersize(A::AbstractSlices) = getinnerdims(size(fused(A)), getsplitmode(A)) +@inline innersize(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = getinnerdims(size(fused(A)), getsplitmode(A)) """ @@ -316,16 +370,69 @@ export innermap innermap(f, obj) = _generic_innermap_impl(f, obj) innermap(f, A::AbstractArray) = map(f, A) innermap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(map, f), A) -innermap(f, A::AbstractSlices) = _generic_innermap_impl(f, A) +innermap(f, A::AbstractSlices{<:AbstractArray}) = _generic_innermap_impl(f, A) function _generic_innermap_impl(f, obj) joined_obj = fused(obj) mapped_joined_obj = map(f, joined_obj) - mapped_obj = splitview(mapped_joined_obj, getsplitmode(obj)) + mapped_obj = splitup(mapped_joined_obj, getsplitmode(obj)) return mapped_obj end +""" + abstract type AbstractPartMode{M,N} <: AbstractSplitMode + +Abstract supertype for array partition modes with `M` inner dimensions and `N` +outer dimensions. + +The mode need not represent a true partition, a partition that discards part +of the original array is allowed. The elements of the partition may also +be reshaped, depending on the mode. + +Use `getsplitmode` to get the split mode of an split array. +""" +abstract type AbstractPartMode{M,N} <: AbstractSplitMode end +export AbstractPartMode + + +""" + vecflattened(A::AbstractArray{T})::AbstractVector{T} + vecflattened(A::AbstractArray{<:AbstractArray})::AbstractVector{T} + +Concatenate nested arrays into a single vector, return non-nested vectors +unchanged. + +If `A` is a nested view of a vector, `vecflattened(A)` should return the +underlying vector in a zero-copy O(1) fashion. So in contrast to +`reduce(vcat, A)` and `mapreduce(vec, vcat, A)`, the result may share +memory with `A`. + +# Implementation + +The default implementations are + +```julia +vecflattened(A::AbstractVector) = A +vecflattened(A::AbstractArray) = vec(A) +vecflattened(A::AbstractVector{<:AbstractVector}) = reduce(vcat, A) +vecflattened(A::AbstractArray{<:AbstractArray}) = mapreduce(vec, vcat, A) +``` + +Specialize `vecflattened` for custom nested array types that can provide a +zero-copy implementation. +""" +function vecflattened end +export vecflattened + +@inline vecflattened(A::AbstractVector) = A +@inline vecflattened(A::AbstractArray) = vec(A) +@inline vecflattened(A::AbstractVector{<:AbstractVector}) = reduce(vcat, A) +@inline vecflattened(A::AbstractVector{<:AbstractArray}) = mapreduce(vec, vcat, A) +@inline vecflattened(A::AbstractArray{<:AbstractVector}) = mapreduce(vec, vcat, A) +@inline vecflattened(A::AbstractArray{<:AbstractArray}) = mapreduce(vec, vcat, A) + + """ deepmap(f, A::AbstractArray) deepmap(f, A::AbstractArray{<:AbstractArray{<:...}}) @@ -339,11 +446,29 @@ export deepmap deepmap(f, obj) = _generic_deepmap_impl(f, obj) deepmap(f, A::AbstractArray) = map(f, A) deepmap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(deepmap, f), A) -deepmap(f, A::AbstractSlices) = _generic_deepmap_impl(f, A) +deepmap(f, A::AbstractSlices{<:AbstractArray}) = _generic_deepmap_impl(f, A) function _generic_deepmap_impl(f, obj) joined_obj = fused(obj) mapped_joined_obj = deepmap(f, joined_obj) - mapped_obj = splitview(mapped_joined_obj, getsplitmode(obj)) + mapped_obj = splitup(mapped_joined_obj, getsplitmode(obj)) return mapped_obj end + + +""" + partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) + partitioned(A::AbstractVector, shapes::AbstractVector{<:Dims}) + +Return a partitioned view of `A`, as a vector of arrays. + +The parts are consecutive, non-overlapping views of `A`, with sizes given by +`lengths` (resulting in a vector of vectors) or `shapes` (resulting in a +vector of arrays). +""" +function partitioned end +export partitioned + +# partitioned methods are defined in vector_of_arrays.jl. + +# ToDo: Add partitioned(A::AbstractVector, n::Integer) and partitioned(A::AbstractVector, shape::Dims) ? diff --git a/src/util.jl b/src/util.jl index 5e8b68b..63269fc 100644 --- a/src/util.jl +++ b/src/util.jl @@ -32,9 +32,12 @@ _convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1 Base.@pure _add_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A + B}() -Base.@pure require_ndims(A::AbstractArray{T,N}, Val_N::Val{N}) where {T,N} = - nothing +Base.@pure _subtract_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A - B}() -Base.@pure require_ndims(A::AbstractArray{T,M}, Val_N::Val{N}) where {T,M,N} = - throw(ArgumentError("Require an array with $N dimensions")) +@inline _require_ndims(::Val{N}, ::Val{N}) where {N} = nothing +function _require_ndims(::Val{N1}, ::Val{N2}) where {N1,N2} + throw(ArgumentError("Require an array with $N2 dimensions, but got an array with $N1 dimensions")) +end + +Base.@pure _val_value(::Val{x}) where x = x diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 791d61c..ba6935f 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -1,12 +1,24 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). +Base.@propagate_inbounds _voa_elem_range(elem_ptr::AbstractVector{<:Integer}, i::Integer) = + elem_ptr[i]:(elem_ptr[i+1] - 1) + +# Element type of a `VectorOfArrays{T,N}` with data of type `VT` and element +# pointers of type `VI`, must match the return type of `getindex` exactly: +@inline function _voa_eltype(::Type{VT}, ::Type{VI}, ::Val{N}) where {VT<:AbstractVector,VI<:AbstractVector{<:Integer},N} + R = Base.promote_op(_voa_elem_range, VI, Int) + SV = Base.promote_op(view, VT, R) + return Base.promote_op(_reshape_dataview, SV, NTuple{N,Int}) +end + + """ - VectorOfArrays{T,N,M} <: AbstractVector{<:AbstractArray{T,N}} + VectorOfArrays{T,N,M,VT,VI,VD,ET<:AbstractArray{T,N}} <: AbstractVector{ET} -An `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may +A `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may differ in size). Internally, `VectorOfArrays` stores all elements of all -arrays in a single flat vector. `M` must equal `N - 1` +arrays in a single flat vector. `M` must equal `N - 1`. The `VectorOfArrays` itself supports `push!`, `unshift!`, etc., but the size of each individual array in the vector is fixed. `resize!` can be used to @@ -38,15 +50,16 @@ and `ArraysOfArrays.no_consistency_checks`. `PartsView` is defined as an type alias: ```julia -`PartsView{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +`PartsView{T,VT,VI,VD,ET} = VectorOfArrays{T,1,0,VT,VI,VD,ET}` ``` """ struct VectorOfArrays{ T, N, M, VT<:AbstractVector{T}, VI<:AbstractVector{<:Integer}, - VD<:AbstractVector{Dims{M}} -} <: AbstractVector{Array{T,N}} + VD<:AbstractVector{Dims{M}}, + ET<:AbstractArray{T,N} +} <: AbstractVector{ET} data::VT elem_ptr::VI kernel_size::VD @@ -57,11 +70,13 @@ struct VectorOfArrays{ elem_ptr = [firstindex(data)] kernel_size = Vector{Dims{M}}() + ET = _voa_eltype(typeof(data), typeof(elem_ptr), Val(N)) new{ T, N, M, typeof(data), typeof(elem_ptr), - typeof(kernel_size) + typeof(kernel_size), + ET }(data, elem_ptr, kernel_size) end @@ -77,7 +92,8 @@ struct VectorOfArrays{ VD<:AbstractVector{Dims{M}} } N = length((ntuple(_ -> 0, Val{M}())..., 0)) - A = new{T,N,M,VT,VI,VD}(data, elem_ptr, kernel_size) + ET = _voa_eltype(VT, VI, Val(N)) + A = new{T,N,M,VT,VI,VD,ET}(data, elem_ptr, kernel_size) checks(A) A end @@ -148,20 +164,16 @@ end Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer) - elem_ptr = A.elem_ptr - - from = elem_ptr[i] - until = elem_ptr[i+1] - to = until - 1 - len = until - from + r = _voa_elem_range(A.elem_ptr, i) + len = length(r) ksize = A.kernel_size[i] klen = prod(ksize) len_p, klen_p = promote(len, klen) sz_lastdim = len == 0 ? len_p : div(len_p, klen_p) - sz = (ksize..., sz_lastdim) + sz = (ksize..., Int(sz_lastdim)) - (from:to, sz) + (r, sz) end @@ -182,6 +194,98 @@ function flatview(A::VectorOfArrays{<:Any,N,M,<:Any,<:SubArray}) where {N,M} view(A.data, A.elem_ptr[begin]:A.elem_ptr[end]-1) end + +""" + struct SplitParts{M,VI,VD} <: AbstractPartMode{M,1} + +The split mode of [`VectorOfArrays`](@ref): a partition of a vector into +consecutive parts of possibly different size, viewed as a vector of +`M`-dimensional arrays. + +Constructor: + +``` +SplitParts( + elem_ptr::AbstractVector{<:Integer}, + kernel_size::AbstractVector{Dims{M-1}} +) +``` + +`elem_ptr` and `kernel_size` equal the equivalent properties of +`VectorOfArrays`. + +See also [`AbstractPartMode`](@ref). +""" +struct SplitParts{ + M, + VI<:AbstractVector{<:Integer}, + VD<:AbstractVector{<:Dims} +} <: AbstractPartMode{M,1} + elem_ptr::VI + kernel_size::VD + + function SplitParts( + elem_ptr::VI, + kernel_size::VD + ) where { + Mk, + VI<:AbstractVector{<:Integer}, + VD<:AbstractVector{Dims{Mk}} + } + M = _val_value(_add_vals(Val(Mk), Val(1))) + new{M,VI,VD}(elem_ptr, kernel_size) + end +end +export SplitParts + +is_memordered_splitmode(::SplitParts) = true + +@inline getsplitmode(A::VectorOfArrays) = SplitParts(A.elem_ptr, A.kernel_size) + +@inline fused(A::VectorOfArrays) = A.data + +function splitup(A::AbstractVector, smode::SplitParts) + VectorOfArrays(A, smode.elem_ptr, smode.kernel_size, simple_consistency_checks) +end + + +function partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) + elem_ptr = _elem_ptr_from_lengths(A, lengths) + kernel_size = similar(lengths, Dims{0}) + fill!(kernel_size, ()) + VectorOfArrays(A, elem_ptr, kernel_size, no_consistency_checks) +end + +function partitioned(A::AbstractVector, shapes::AbstractVector{Dims{N}}) where {N} + elem_ptr = _elem_ptr_from_lengths(A, prod.(shapes)) + kernel_size = Base.front.(shapes) + VectorOfArrays(A, elem_ptr, kernel_size, no_consistency_checks) +end + +function _elem_ptr_from_lengths(A::AbstractVector, lengths::AbstractVector{<:Integer}) + elem_ptr = similar(lengths, Int, length(lengths) + 1) + i = firstindex(elem_ptr) + elem_ptr[i] = firstindex(A) + for l in lengths + l >= 0 || throw(ArgumentError("Part lengths must not be negative")) + elem_ptr[i + 1] = elem_ptr[i] + l + i += 1 + end + last(elem_ptr) - 1 <= lastindex(A) || throw(ArgumentError("Sum of part lengths exceeds length of data vector")) + return elem_ptr +end + + +function vecflattened(A::VectorOfArrays) + ep = A.elem_ptr + view(A.data, first(ep):(last(ep) - 1)) +end + +# Fast paths, must return independent arrays, unlike `vecflattened`: +Base.mapreduce(::typeof(vec), ::typeof(vcat), A::VectorOfArrays) = copy(vecflattened(A)) +Base.reduce(::typeof(vcat), A::VectorOfArrays{T,1}) where {T} = copy(vecflattened(A)) + + Base.size(A::VectorOfArrays) = size(A.kernel_size) Base.IndexStyle(::Type{<:VectorOfArrays}) = IndexLinear() @@ -388,7 +492,7 @@ function Base.push!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}) where {T,N,U} end -function Base.empty(A::VectorOfArrays{T,N}, ::Type{<:DenseArray{U,N}}) where {T,N,U} +function Base.empty(A::VectorOfArrays{T,N}, ::Type{<:AbstractArray{U,N}}) where {T,N,U} empty_data = empty(A.data, U) empty_elem_ptr = push!(empty(A.elem_ptr), firstindex(empty_data)) empty_kernel_size = empty(A.kernel_size) @@ -421,7 +525,10 @@ Base.Broadcast.broadcasted(::typeof(identity), A::VectorOfArrays) = A """ - PartsView{T,...} = VectorOfArrays{T,1,...} + PartsView{T,...} = VectorOfArrays{T,1,0,...} + +A vector of vectors (that may differ in length), stored in contiguous, +partitioned form. See [`VectorOfArrays`](@ref) for details. Constructors: @@ -433,16 +540,17 @@ PartsView( data::AbstractVector, elem_ptr::AbstractVector{<:Integer}, checks::Function = full_consistency_checks ) - -See also [VectorOfArrays](@ref). ``` + +See also [`VectorOfArrays`](@ref). """ const PartsView{ T, VT<:AbstractVector{T}, VI<:AbstractVector{<:Integer}, - VD<:AbstractVector{Dims{0}} -} = VectorOfArrays{T,1,0,VT,VI,VD} + VD<:AbstractVector{Dims{0}}, + ET<:AbstractVector{T} +} = VectorOfArrays{T,1,0,VT,VI,VD,ET} export PartsView @@ -577,9 +685,4 @@ end # Deprecated: -const VectorOfArrays{T,N} = PartsView{T,N} -export VectorOfArrays - - -const VectorOfVectors{T} = VoVPartsView{T} -export VectorOfVectors +Base.@deprecate_binding VectorOfVectors PartsView diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 0c0e173..22c3127 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -32,17 +32,16 @@ using StatsBase: cov2cor @inline function test_from_flat(::Type{TT}, ::Type{RT}, Val_L::Val{L}) where { - TT<:SlicedView, - RT<:SlicedView, + TT<:ArrayOfSimilarArrays, + RT<:ArrayOfSimilarArrays, L } @testset "$TT from Array{Float64,$L}" begin A = rand_flat_array(Val_L) @test typeof(@inferred TT(A)) <: RT - @test typeof(@inferred convert(TT, A)) <: RT AosA = TT(A) - @test typeof(AosA) == typeof(convert(TT, A)) + @test typeof(AosA) == typeof(TT(A)) @test @inferred(eltype(AosA)) == typeof(AosA[1]) @test @inferred(parent(AosA)) === @inferred(flatview(AosA)) if eltype(eltype(AosA)) == eltype(A) @@ -50,67 +49,72 @@ using StatsBase: cov2cor else @test @inferred(flatview(AosA)) ≈ A end - @test @inferred(stack(AosA)) === flatview(AosA) + @test @inferred(stacked(AosA)) === flatview(AosA) + @test @inferred(stack(AosA)) == flatview(AosA) + @test stack(AosA) !== flatview(AosA) + + @test_deprecated convert(TT, A) == TT(A) end end @inline function test_from_nested(::Type{TT}, ::Type{RT}, Val_M::Val{M}, Val_N::Val{N}) where { - TT<:SlicedView, - RT<:SlicedView, + TT<:ArrayOfSimilarArrays, + RT<:ArrayOfSimilarArrays, M, N } @testset "$TT from Array{Array{Float64,$M},$N}" begin A = rand_nested_similar_arrays(Val_M, Val_N) - A2_ctor = @inferred TT(A) - U = eltype(flatview(A2_ctor)) + A2_conv = @inferred convert(TT, A) + U = eltype(flatview(A2_conv)) A_U = Array{Array{U,M},N}(A) - @test typeof(A2_ctor) <: RT - @test A2_ctor == A_U + @test typeof(A2_conv) <: RT + @test A2_conv == A_U A2_conv = @inferred convert(TT, A) @test typeof(A2_conv) <: RT - @test A2_conv == A2_ctor - U = eltype(flatview(A2_ctor)) - A3 = @inferred Array(A2_ctor) + U = eltype(flatview(A2_conv)) + A3 = @inferred Array(A2_conv) @test typeof(A3) == Array{Array{U,M},N} @test A3 == A_U + + @test_deprecated TT(A) == A2_conv end end @testset "construct/convert from flat array" begin - test_from_flat(SlicedView{Float64,2,3}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(5)) - test_from_flat(SlicedView{Float64,2}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) - test_from_flat(SlicedView{Float32,2,3}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(5)) - test_from_flat(SlicedView{Float32,2}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) - test_rrule(SlicedView{Float64,2,2}, rand(2,3,4,5)) + test_rrule(ArrayOfSimilarArrays{Float64,2,2}, rand(2,3,4,5)) end @testset "construct/convert from nested arrays" begin - test_from_nested(SlicedView, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(SlicedView{Float64,2,3}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(SlicedView{Float64}, SlicedView{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(SlicedView{Float32}, SlicedView{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) - test_from_nested(SlicedView, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(SlicedView, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(SlicedView, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) + test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) @@ -129,12 +133,12 @@ using StatsBase: cov2cor r = @inferred(rand(5,5)) @test @inferred(flatview(ArrayOfSimilarVectors(r))) == r - test_rrule(SlicedView{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) + test_rrule(ArrayOfSimilarArrays{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) end @testset "AbstractSlices interface" begin - @test @inferred(SlicedView{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices - let A = SlicedView{Float32,2,3}(rand(3,4,5,6,7)) + @test @inferred(ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices + let A = ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7)) @test @inferred(A[2, 3, 4]) isa AbstractArray{Float32,2} ref_ET = typeof(A[2, 3, 4]) @test A isa AbstractSlices{ref_ET, 3} @@ -150,8 +154,8 @@ using StatsBase: cov2cor EA_ref2 = rand_flat_array(Val(3)) EA1 = ElasticArray{Float64, 3}(EA_ref1) EA2 = ElasticArray{Float64, 3}(EA_ref2) - AEA1 = SlicedView{Float64, 3}(EA1) - AEA2 = SlicedView{Float64, 3}(EA2) + AEA1 = ArrayOfSimilarArrays{Float64, 3}(EA1) + AEA2 = ArrayOfSimilarArrays{Float64, 3}(EA2) AEA1_ref = copy(AEA1) AEA2_ref = copy(AEA2) append!(AEA1, AEA2) @@ -161,8 +165,8 @@ using StatsBase: cov2cor @test @inferred(reverse(AEA2.data, dims=3)[:,:,i]) == @inferred(AEA1.data[:,:,N+1-i]) end - A1 = SlicedView{Float64,1}(rand_flat_array(Val(1))) - A2 = SlicedView{Float64,1}(rand_flat_array(Val(1))) + A1 = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) + A2 = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) A1_data = copy(A1.data) A2_data = copy(A2.data) append!(A1, A2) @@ -173,14 +177,14 @@ using StatsBase: cov2cor end @testset "similar and copyto!" begin - A = SlicedView{Float64,1}(rand_flat_array(Val(1))) + A = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) @test (@inferred copyto!((@inferred similar(A)), A)) == A - A = SlicedView{Float64,2}(rand_flat_array(Val(5))) + A = ArrayOfSimilarArrays{Float64,2}(rand_flat_array(Val(5))) @test (@inferred copyto!((@inferred similar(A)), A)) == A A_data = rand_flat_array(Val(4)) - A = SlicedView{Float64, 2}(A_data) + A = ArrayOfSimilarArrays{Float64, 2}(A_data) A_similar = similar(A, Array{Float64, 2}, size(A)) @test @inferred(size(A)) == @inferred(size(A_similar)) @test @inferred(size(A.data)) == @inferred(size(A_similar.data)) @@ -191,14 +195,14 @@ using StatsBase: cov2cor @testset "adapt" begin A_flat = rand(2,3,4,5,6) - A_nested = nestedview(A_flat, 2) + A_nested = sliced(A_flat, 2) @test @inferred(adapt(identity, A_nested)) == A_nested @test typeof(adapt(identity, A_nested)) == typeof(A_nested) end @testset "deepcopy" begin - A = SlicedView{Float64,1}(rand_flat_array(Val(1))) + A = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1))) @test (@inferred deepcopy(A)) == A @test typeof(deepcopy(A)) == typeof(A) end @@ -206,14 +210,16 @@ using StatsBase: cov2cor @testset "flatview" begin A = rand_nested_similar_arrays(Val(3), Val(2)) - B = SlicedView(A) - @inferred(flatview(B))[:] == collect(flatview(A)) + B = convert(ArrayOfSimilarArrays, A) + @test @inferred(flatview(B)) === B.data + @test flatview(B) == stacked(A) + @test_throws ArgumentError flatview(A) end @testset "empty" begin A = [rand(2,3), rand(2,3), rand(2,3)] - B = SlicedView(A) + B = ArrayOfSimilarArrays(A) @test typeof(@inferred empty(B)) == typeof(B) @test empty(A) == empty(B) @@ -227,10 +233,10 @@ using StatsBase: cov2cor @testset "stats" begin VV = [rand(3) for i in 1:10] - VV_aosa = SlicedView(VV) + VV_aosa = ArrayOfSimilarArrays(VV) VA = [rand(2,3,3) for i in 1:10] - VA_aosa = SlicedView(VA) + VA_aosa = ArrayOfSimilarArrays(VA) array_cmp(A, B) = (A ≈ B) && (size(A) == size(B)) @@ -264,13 +270,13 @@ using StatsBase: cov2cor @testset "examples" begin A_flat = rand(2,3,4,5,6) - A_nested = nestedview(A_flat, 2) + A_nested = sliced(A_flat, 2) @test A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T @test flatview(A_nested) === A_flat A_flat = rand(4,4) - A_nested = @inferred(nestedview(A_flat)) + A_nested = @inferred(sliced(A_flat)) @test A_nested.data == A_flat @test @inferred(size(A_nested))[1] == @inferred(size(A_flat))[1] @@ -279,11 +285,11 @@ using StatsBase: cov2cor # ------------------------------------------------------------------- A_flat = rand(2,3,4,5) - ASA = @inferred(SlicedView{Float64,2,2}(A_flat)) + ASA = @inferred(ArrayOfSimilarArrays{Float64,2,2}(A_flat)) @test ASA.data == A_flat # ------------------------------------------------------------------- - A_nested = nestedview(ElasticArray{Float64}(undef, 2, 3, 0), 2) + A_nested = sliced(ElasticArray{Float64}(undef, 2, 3, 0), 2) A_nested_copy = deepcopy(A_nested) for i in 1:4 @@ -317,7 +323,7 @@ using StatsBase: cov2cor r = vcat(r1,r2,r3,r4) VSV = VectorOfSimilarVectors(r) VSA = VectorOfSimilarArrays(r) - ASA = SlicedView([r1,r2,r3,r4]) + ASA = ArrayOfSimilarArrays([r1,r2,r3,r4]) f = x -> x.*2 @@ -335,7 +341,7 @@ using StatsBase: cov2cor @testset "map and broadcast" begin A_flat = rand(2,3,4,5,6) - A = nestedview(A_flat, 2) + A = sliced(A_flat, 2) for do_map in (map, broadcast) @test @inferred(do_map(identity, A)) === A diff --git a/test/base_slices.jl b/test/base_slices.jl index 3abd618..ac8a54b 100644 --- a/test/base_slices.jl +++ b/test/base_slices.jl @@ -21,6 +21,13 @@ include("testdefs.jl") @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes2))) == (2,4) @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes2))) == (3,1,5) + # Non-involutive outer dimension order (a 3-cycle), regression test for + # getouterdims applying the inverse of the slicemap permutation: + Aes3 = eachslice(A_orig; dims = (2,3,1)) + @test @inferred(getsplitmode(Aes3)) isa BaseSlicing{2,3,Tuple{Int,Int,Int,Colon,Colon}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes3))) == (4,5) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes3))) == (2,3,1) + Aec = eachcol(A_orig_mat) @test @inferred(getsplitmode(Aec)) isa BaseSlicing{1,1,Tuple{Colon,Int}} @test @inferred(getinnerdims((1,2), getsplitmode(Aec))) == (1,) @@ -33,6 +40,7 @@ include("testdefs.jl") test_api(Aes1, Array(Aes1), A_orig) test_api(Aes2, Array(Aes2), A_orig) + test_api(Aes3, Array(Aes3), A_orig) test_api(Aec, Array(Aec), A_orig_mat) test_api(Aer, Array(Aer), A_orig_mat) end diff --git a/test/broadcasting.jl b/test/broadcasting.jl index cacf3c9..19cce78 100644 --- a/test/broadcasting.jl +++ b/test/broadcasting.jl @@ -27,7 +27,7 @@ using Test end end - let A = SlicedView(ref_AosA1(Float32, 100)) + let A = ArrayOfSimilarArrays(ref_AosA1(Float32, 100)) refA = Array(A) for Idxs in [ @@ -45,7 +45,7 @@ using Test ] refA = Array(A) - @test @inferred(broadcast(getindex, A, Idxs...)) isa SlicedView{eltype(eltype(A))} + @test @inferred(broadcast(getindex, A, Idxs...)) isa ArrayOfSimilarArrays{eltype(eltype(A))} @test getindex.(A, Idxs...) == getindex.(refA, Idxs...) end end @@ -54,7 +54,7 @@ using Test @testset "findall" begin for A in [ VectorOfArrays(ref_VoA1(Bool, 100)), - SlicedView(ref_AosA1(Bool, 100)) + ArrayOfSimilarArrays(ref_AosA1(Bool, 100)) ] refA = Array(A) diff --git a/test/static_arrays.jl b/test/static_arrays.jl index ee136a4..77fd454 100644 --- a/test/static_arrays.jl +++ b/test/static_arrays.jl @@ -7,11 +7,12 @@ using StaticArrays @testset "StaticArray Extension" begin - @testset "flatview and nestedview" begin + @testset "flatview and sliced" begin A = [(@SArray randn(3, 2, 4)) for i in 1:2, j in 1:2] - @test @inferred(nestedview(flatview(A), Val(3))) == A + @test @inferred(sliced(flatview(A), Val(3))) == A B = rand(3, 2, 4) - @test @inferred(nestedview(flatview(B), SVector{3})) == @inferred(nestedview(B, Val(1))) + @test @inferred(sliced(flatview(B), SVector{3})) == @inferred(sliced(B, Val(1))) + @test_deprecated nestedview(flatview(B), SVector{3}) end end diff --git a/test/testdefs.jl b/test/testdefs.jl index 69ab712..099a4d6 100644 --- a/test/testdefs.jl +++ b/test/testdefs.jl @@ -2,7 +2,7 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent -#if !isdefined(Main, :test_api) +if !isdefined(Main, :test_api) maptest_f(x::Number) = x^2 maptest_f(x::AbstractArray{<:Number}) = sum(x)^2 maptest_f(x::AbstractArray) = length(x)^2 @@ -14,7 +14,7 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent @test @inferred(f_pullback(dy)) == ntuple(_ -> NoTangent(), Val(N+1)) end - function test_joinedview_rrule(A) + function test_fused_rrule(A) @testset "Test fused rrule for $(nameof(typeof(A)))" begin A_joined, f_pullback = @inferred(rrule(fused, A)) @test A_joined == fused(A) @@ -26,9 +26,9 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent end function test_splitview_rrule(A, smode) - @testset "Test splitview rrule for $(nameof(typeof(A)))" begin - A_split, f_pullback = @inferred(rrule(splitview, A, smode)) - @test A_split == splitview(A, smode) + @testset "Test splitup rrule for $(nameof(typeof(A)))" begin + A_split, f_pullback = @inferred(rrule(splitup, A, smode)) + @test A_split == splitup(A, smode) dy = Thunk(Returns(A_split)) tangents = @inferred(f_pullback(dy)) @test tangents isa Tuple{NoTangent,<:AbstractThunk,NoTangent} @@ -104,15 +104,18 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent @test Array(stack(A)) == A_array_stacked + @test @inferred(unstackmode(A)) isa AbstractSlicingMode{M,N} + umode = unstackmode(A) + @test @inferred(splitup(stacked(A), umode)) == A + if is_memordered_splitmode(smode) - if A isa Slices - # stack(A) never returns parent for Slices, even if possible: - @test @inferred(stack(A)) == A_unsplit_ref - else - @test @inferred(stack(A)) === A_unsplit_ref - end + # stack(A) must return an independent array, even if the + # parent array could be returned as-is: + @test @inferred(stack(A)) == A_unsplit_ref + @test stack(A) !== A_unsplit_ref @test @inferred(stacked(A)) === A_unsplit_ref @test @inferred(flatview(A)) === A_unsplit_ref + @test stacked(splitup(stacked(A), umode)) === stacked(A) else @test Array(@inferred(stack(A))) == A_array_stacked @test Array(@inferred(stacked(A))) == A_array_stacked @@ -139,18 +142,18 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent @test @inferred(is_memordered_splitmode(smode)) == false @test_throws ArgumentError fused(A) @test_throws ArgumentError flatview(A) - @test_throws ArgumentError splitview(A_unsplit_ref, smode) + @test_throws ArgumentError splitup(A_unsplit_ref, smode) @test_throws ArgumentError getinnerdims(dimstpl, smode) @test_throws ArgumentError getouterdims(dimstpl, smode) elseif smode isa NonSplitMode @test @inferred(is_memordered_splitmode(smode)) == true @test @inferred(fused(A)) === A @test @inferred(flatview(A)) === A - @test @inferred(splitview(A, smode)) === A + @test @inferred(splitup(A, smode)) === A @test @inferred(getinnerdims(dimstpl, smode)) == () @test @inferred(getouterdims(dimstpl, smode)) == dimstpl - test_joinedview_rrule(A) + test_fused_rrule(A) test_splitview_rrule(A, smode) else if A isa Slices @@ -159,12 +162,12 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent @test @inferred(fused(A)) == A_unsplit_ref A_unsplit = fused(A) @test typeof(A_unsplit) == typeof(A_unsplit_ref) - @test typeof(splitview(A_unsplit, smode)) == typeof(A) - @test splitview(A_unsplit, smode) == A + @test typeof(splitup(A_unsplit, smode)) == typeof(A) + @test splitup(A_unsplit, smode) == A - test_joinedview_rrule(A) + test_fused_rrule(A) test_splitview_rrule(A_unsplit, smode) end end end -#end +end diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 6021347..617d3ae 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -8,6 +8,8 @@ using Adapt using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr +include("testdefs.jl") + @testset "vector_of_arrays" begin ref_flatview(A::AbstractVector{<:AbstractArray}) = vcat(map(vec, Array(A))...) @@ -88,17 +90,18 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test @inferred(vcat(B1, B2)) isa VectorOfArrays @test vcat(B1, B2) == vcat(A1, A2) - @test eltype(vcat(B1, B2)) == Array{Float32,3} + @test eltype(vcat(B1, B2)) <: AbstractArray{Float32,3} + @test eltype(vcat(B1, B2)) == eltype(B1) full_consistency_checks(vcat(B1, B2)) @test @inferred(vcat(B1, B3)) isa VectorOfArrays @test vcat(B1, B3) == vcat(A1, A3) - @test eltype(vcat(B1, B3)) == Array{Float32,3} + @test eltype(vcat(B1, B3)) <: AbstractArray{Float32,3} full_consistency_checks(vcat(B1, B3)) @test @inferred(vcat(B1, B2, B3, B4)) isa VectorOfArrays @test vcat(B1, B2, B3, B4) == vcat(A1, A2, A3, A4) - @test eltype(vcat(B1, B2, B3, B4)) == Array{Float64,3} + @test eltype(vcat(B1, B2, B3, B4)) <: AbstractArray{Float64,3} full_consistency_checks(vcat(B1, B2, B3, B4)) B1_copy = @inferred(copy(B1)); B3_copy = @inferred(copy(B3)) @@ -107,6 +110,68 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr end + @testset "split mode API" begin + B1 = VectorOfArrays(ref_AoA1(Float32, 5)) + B1e = VectorOfArrays(ref_AoA1(Float32, 0)) + B3 = VectorOfArrays(ref_AoA3(Float32, 3)) + Bu = VectorOfArrays([rand(Float32, 2, 3) for i in 1:4]) + + for B in [B1, B1e, B3, Bu] + @test @inferred(getsplitmode(B)) isa AbstractPartMode{ndims(eltype(B)),1} + @test @inferred(fused(B)) === B.data + @test @inferred(splitup(fused(B), getsplitmode(B))) == B + @test typeof(splitup(fused(B), getsplitmode(B))) == typeof(B) + test_api(B, Array(B), B.data) + end + + # Uniform element size, so stackable: + @test @inferred(stacked(Bu)) == stack(Array(Bu)) + @test @inferred(splitup(stacked(Bu), unstackmode(Bu))) == Bu + + # flatview on view-backed VectorOfArrays returns a view of only the + # covered data range, in a type-stable fashion: + B3_view = view(B3, 2:3) + @test @inferred(flatview(B3_view)) isa SubArray + @test flatview(B3_view) == B3.data[B3.elem_ptr[2]:(B3.elem_ptr[4] - 1)] + @test @inferred(vecflattened(B3_view)) == flatview(B3_view) + @test @inferred(flatview(B3)) === B3.data + end + + + @testset "partitioned" begin + x = collect(1:10) + + p = @inferred(partitioned(x, [2, 3, 5])) + @test p isa PartsView{Int} + @test p == [[1, 2], [3, 4, 5], [6, 7, 8, 9, 10]] + @test @inferred(fused(p)) === x + @test flatview(p) === x + test_api(p, Array(p), x) + + # vecflattened may share memory, reduce/mapreduce must not: + @test @inferred(vecflattened(p)) == x + @test parent(vecflattened(p)) === x + @test @inferred(reduce(vcat, p)) == x + @test reduce(vcat, p) !== x + @test typeof(reduce(vcat, p)) == typeof(x) + @test @inferred(mapreduce(vec, vcat, p)) == x + @test mapreduce(vec, vcat, p) !== x + + # Partial partitions are allowed: + p_partial = @inferred(partitioned(x, [2, 3])) + @test p_partial == [[1, 2], [3, 4, 5]] + + @test_throws ArgumentError partitioned(x, [2, 3, 6]) + @test_throws ArgumentError partitioned(x, [2, -1, 5]) + + p2 = @inferred(partitioned(x, [(1, 2), (2, 2)])) + @test p2 isa VectorOfArrays{Int,2} + @test p2[1] == [1 2] + @test p2[2] == [3 5; 4 6] + @test @inferred(fused(p2)) === x + end + + @testset "indexing" begin V1 = @inferred(VectorOfArrays(ref_AoA3(Float32, 3))) V2 = @inferred(VectorOfArrays(ref_AoA3(Float32, 3))) From caa8e71fceb86d74896bd249eaf4d08a5f79aca6 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 01:13:56 +0200 Subject: [PATCH 21/45] Fix spelling and grammar --- README.md | 2 +- docs/src/index.md | 2 +- src/array_of_similar_arrays.jl | 2 +- src/functions.jl | 22 +++++++++++----------- src/vector_of_arrays.jl | 8 ++++---- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index efd24e9..e0b70b5 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: `ArrayOfSimilarArrays` and `VectorOfArrays`. -An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. +An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user via `flatview()`. ## Documentation diff --git a/docs/src/index.md b/docs/src/index.md index 85ada8f..1329754 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -94,7 +94,7 @@ VA_flat = flatview(VA) VA_flat isa Vector{Float64} ``` -Calling `getindex` on `A_nested` returns a view into `A_flat`: +Calling `getindex` on `VA` returns a view into `VA_flat`: ```julia VA_flat = flatview(VA) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 2ff7a0f..894fb48 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -76,7 +76,7 @@ The following type aliases are defined: * `VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}` `VectorOfSimilarArrays` supports `push!()`, etc., provided the underlying -array supports resizing of it's last dimension (e.g. an `ElasticArray`). +array supports resizing of its last dimension (e.g. an `ElasticArray`). The nested array can also be created using the function [`sliced`](@ref) and the wrapped flat array can be accessed using [`flatview`](@ref) diff --git a/src/functions.jl b/src/functions.jl index 694a4b0..b27812c 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -36,10 +36,10 @@ export AbstractSplitMode """ is_memordered_splitmode(smode::AbstractSplitMode)::Bool -Check if `smode` splits in memory-order. - +Check whether `smode` splits in memory order. + If true, inner arrays are stored contiguously in memory in Julia-native -dimension order, and the same is true for the outer dimensions (no dimention +dimension order, and the same is true for the outer dimensions (no dimension reordering). If true, `flatview` and `fused` are equivalent. @@ -69,7 +69,7 @@ unknown way, e.g. nested arrays of type `Array{<:Array}`. Since the split parts may be (typically are) non-contiguous in memory, this split mode does not allow for `fused` or `flatview`. It is also not -inferrable if the split object should be interpreted as a sliced array, +inferable whether the split object should be interpreted as a sliced array, a ragged array, or something else. Constructor: `UnknownSplitMode{T}()` @@ -131,7 +131,7 @@ View array `A` in unsplit form. the same type as `A` if at all possible, except if `getsplitmode(A)` is an `UnknownSplitMode`. -If `A` is not a nested array return `A` itself. If `A` is a split array, +If `A` is not a nested array, return `A` itself. If `A` is a split array, return the original unsplit array. If `is_memordered_splitmode(getsplitmode(A))` is true, `fused(A)` is @@ -160,7 +160,7 @@ end flatview(A::AbstractArray{<:AbstractArray}) View array `A` in a flattened form, with inner dimensions first. The shape of -the flattened form will depend on the type of `A`. If the `A` is not a +the flattened form will depend on the type of `A`. If `A` is not a nested array, the return value is `A` itself. Only specific types of nested arrays are supported. @@ -198,7 +198,7 @@ end Abstract supertype for array slicing modes with `M` inner dimensions and `N` outer dimensions. -Use `getsplitmode` to get the split mode of an split array. +Use `getsplitmode` to get the split mode of a split array. """ abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode end export AbstractSlicingMode @@ -312,7 +312,7 @@ end sliced(A::AbstractArray{T,M+N}, Val(M)) sliced(A::AbstractArray{T,M+N}, M::Integer) -Return a sliced view of `A`, using the columns or the the first `M` +Return a sliced view of `A`, using the columns or the first `M` dimensions as inner dimensions. """ function sliced end @@ -361,7 +361,7 @@ end innermap(f, A::AbstractArray) innermap(f, A::AbstractArray{<:AbstractArray}) -Nested `map` at depth 2. Equivalent to `map(X -> map(f, X) A)` for arrays +Nested `map` at depth 2. Equivalent to `map(X -> map(f, X), A)` for arrays of arrays, otherwise equivalent to `Base.map`. """ function innermap end @@ -390,7 +390,7 @@ The mode need not represent a true partition, a partition that discards part of the original array is allowed. The elements of the partition may also be reshaped, depending on the mode. -Use `getsplitmode` to get the split mode of an split array. +Use `getsplitmode` to get the split mode of a split array. """ abstract type AbstractPartMode{M,N} <: AbstractSplitMode end export AbstractPartMode @@ -438,7 +438,7 @@ export vecflattened deepmap(f, A::AbstractArray{<:AbstractArray{<:...}}) Applies `map` at the deepest layer of nested arrays. If `A` is not -a nested array, `deepmap` behaves identical to `Base.map`. +a nested array, `deepmap` behaves identically to `Base.map`. """ function deepmap end export deepmap diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index ba6935f..5568010 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -20,7 +20,7 @@ A `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may differ in size). Internally, `VectorOfArrays` stores all elements of all arrays in a single flat vector. `M` must equal `N - 1`. -The `VectorOfArrays` itself supports `push!`, `unshift!`, etc., but the size +The `VectorOfArrays` itself supports `push!`, `append!`, etc., but the size of each individual array in the vector is fixed. `resize!` can be used to shrink, but not to grow, as the size of the additional element arrays in the vector would be unknown. However, memory space for up to `n` arrays with a @@ -120,7 +120,7 @@ Base.convert(VA::Type{VectorOfArrays}, A::AbstractVector{AbstractArray{T,N}}) wh Returns the internal element pointer vector of `A`. -Do *not* change modify the returned vector in any way, as this would break the +Do *not* modify the returned vector in any way, as this would break the inner consistency of `A`. Use with care, see [`element_ptr`](@ref) for a safe version of this function. @@ -586,7 +586,7 @@ Example: elem_ptr = consgrouped_ptrs(A) first.(PartsView(A, elem_ptr)) == [1, 2, 3, 2] ``` -consgrouped_ptrs + Typically, `elem_ptr` will be used to apply the computed grouping to other data: @@ -623,7 +623,7 @@ end Compute a grouping of equal consecutive elements on `source` via [`consgrouped_ptrs`](@ref) and apply the grouping to target, resp. each -element of `target`. `target` may be an vector or a named or unnamed tuple of +element of `target`. `target` may be a vector or a named or unnamed tuple of vectors. The result is a `PartsView`, resp. a tuple of such. Example: From 455bde7a13a3b87fec95e92131d8f9f6a0100852 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 01:15:46 +0200 Subject: [PATCH 22/45] FIX: rename _convert_elype to _convert_eltype --- src/array_of_similar_arrays.jl | 2 +- src/broadcasting.jl | 2 +- src/util.jl | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 894fb48..3e0c385 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -108,7 +108,7 @@ export ArrayOfSimilarArrays function ArrayOfSimilarArrays{T,M,N}(orig_data::AbstractArray{U,L}) where {T,M,N,U,L} - conv_data = _convert_elype(T, orig_data)::AbstractArray{T,L} + conv_data = _convert_eltype(T, orig_data)::AbstractArray{T,L} return ArrayOfSimilarArrays{T,M,N}(conv_data) end diff --git a/src/broadcasting.jl b/src/broadcasting.jl index e76cad2..9da047e 100644 --- a/src/broadcasting.jl +++ b/src/broadcasting.jl @@ -35,7 +35,7 @@ _noreindex_view(A::AbstractArray{T,N}, ::Vararg{Colon,N}) where {T,N} = A _generic_size(A) = size(A) _generic_size(tpl::Tuple) = (length(tpl),) -Base.Base.@propagate_inbounds function _to_indices(A, idxs) +Base.@propagate_inbounds function _to_indices(A, idxs) new_idxs = Base.to_indices(A, idxs) @boundscheck Base.checkbounds(A, new_idxs...) return new_idxs diff --git a/src/util.jl b/src/util.jl index 63269fc..b3034d5 100644 --- a/src/util.jl +++ b/src/util.jl @@ -25,9 +25,9 @@ Base.@propagate_inbounds swap_front_back_tuple(x::NTuple{N,Any}, ::Val{M}) where (back_tuple(x, Val{N - M}())..., front_tuple(x, Val{M}())...) -_convert_elype(::Type{T}, A::AbstractArray{T}) where {T} = A +_convert_eltype(::Type{T}, A::AbstractArray{T}) where {T} = A -_convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1(convert, T), A) +_convert_eltype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1(convert, T), A) Base.@pure _add_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A + B}() From 3af39379613c66e80744e5e1ad3d6ed57209ba1e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 01:54:37 +0200 Subject: [PATCH 23/45] Fix equality semantics of VectorOfArrays == compared internal fields directly, so vectors of arrays with equal elements but different data layout (offsets, unused data regions) compared unequal, contradicting both elementwise semantics and isequal. Compare element sizes and covered data instead, and specialize isequal consistently. Co-Authored-By: Claude Fable 5 --- src/vector_of_arrays.jl | 28 ++++++++++++++++++++++++++-- test/vector_of_arrays.jl | 30 ++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 5568010..e3573fb 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -177,9 +177,33 @@ Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer end +# Equality must be equivalent to elementwise comparison, but can be checked +# much more efficiently via the flat data. The underlying data of A and B may +# be offset relative to their element pointers, so only element sizes and the +# data regions covered by the elements may be compared: + +function _elem_lengths_equal(elem_ptr_A::AbstractVector{<:Integer}, elem_ptr_B::AbstractVector{<:Integer}) + length(elem_ptr_A) == length(elem_ptr_B) || return false + iA, iB = firstindex(elem_ptr_A), firstindex(elem_ptr_B) + @inbounds for k in 0:(length(elem_ptr_A) - 2) + elem_ptr_A[iA+k+1] - elem_ptr_A[iA+k] == elem_ptr_B[iB+k+1] - elem_ptr_B[iB+k] || return false + end + return true +end + import Base.== -(==)(A::VectorOfArrays, B::VectorOfArrays) = - A.data == B.data && A.elem_ptr == B.elem_ptr && A.kernel_size == B.kernel_size +function ==(A::VectorOfArrays{<:Any,N}, B::VectorOfArrays{<:Any,N}) where {N} + axes(A) == axes(B) || return false + A.kernel_size == B.kernel_size || return false + _elem_lengths_equal(A.elem_ptr, B.elem_ptr) || return false + return vecflattened(A) == vecflattened(B) +end + +function Base.isequal(A::VectorOfArrays{<:Any,N}, B::VectorOfArrays{<:Any,N}) where {N} + axes(A) == axes(B) && isequal(A.kernel_size, B.kernel_size) && + _elem_lengths_equal(A.elem_ptr, B.elem_ptr) && + isequal(vecflattened(A), vecflattened(B)) +end """ flatview(A::VectorOfArrays{T})::Vector{T} diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 617d3ae..3331e11 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -110,6 +110,36 @@ include("testdefs.jl") end + @testset "equality" begin + # == and isequal must be equivalent to elementwise comparison, + # independent of the underlying data layout: + p = partitioned(collect(1:10), [2, 3]) # unused data after covered range + q = VectorOfArrays([[1, 2], [3, 4, 5]]) + @test collect(p) == collect(q) + @test p == q + @test q == p + @test isequal(p, q) + + @test p != VectorOfArrays([[1, 2], [3, 4, 6]]) + @test p != VectorOfArrays([[1, 2, 3], [4, 5]]) + @test p != VectorOfArrays([[1, 2], [3, 4, 5], [6]]) + + # Equal-length parts with different shapes are not equal: + r1 = VectorOfArrays([[1 2; 3 4]]) + r2 = VectorOfArrays([[1 3; 2 4]]) + r3 = partitioned(collect([1, 3, 2, 4]), [(2, 2)]) + @test r1 != r2 + @test r1 == permutedims.(r2) + @test r1 == r3 && r2 != r3 + + # missing propagates through ==, but not isequal: + m1 = VectorOfArrays([[1, missing], [3]]) + m2 = VectorOfArrays([[1, missing], [3]]) + @test ismissing(m1 == m2) + @test isequal(m1, m2) + end + + @testset "split mode API" begin B1 = VectorOfArrays(ref_AoA1(Float32, 5)) B1e = VectorOfArrays(ref_AoA1(Float32, 0)) From e301be167b0397ed801fc8e2fdb460bb77d888c7 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 01:56:32 +0200 Subject: [PATCH 24/45] Fix element type and outer dims in similar for ArrayOfSimilarArrays similar(A, Array{U,M}, dims) allocated data with element type U but then wrapped it as ArrayOfSimilarArrays{T,...}, converting back to the old element type. It also hard-coded the old number of outer dimensions instead of using length(dims). Co-Authored-By: Claude Fable 5 --- src/array_of_similar_arrays.jl | 2 +- test/array_of_similar_arrays.jl | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 3e0c385..e1fa6dc 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -189,7 +189,7 @@ function Base.similar(A::ArrayOfSimilarArrays{T,M,N}, ::Type{<:AbstractArray{U}} data = A.data size_inner = front_tuple(size(data), Val{M}()) # ToDo: Don't use similar if data is an ElasticArray? - ArrayOfSimilarArrays{T,M,N}(similar(data, U, size_inner..., dims...)) + ArrayOfSimilarArrays{U,M,length(dims)}(similar(data, U, size_inner..., dims...)) end diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 22c3127..f02c42f 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -189,6 +189,14 @@ using StatsBase: cov2cor @test @inferred(size(A)) == @inferred(size(A_similar)) @test @inferred(size(A.data)) == @inferred(size(A_similar.data)) @test typeof(A_similar.data) == typeof(A_data) + + # similar must respect the requested element type and outer dims: + A_similar32 = @inferred similar(A, Array{Float32, 2}, size(A)) + @test eltype(eltype(A_similar32)) == Float32 + @test eltype(A_similar32.data) == Float32 + A_similar_1d = @inferred similar(A, Array{Float64, 2}, (3,)) + @test size(A_similar_1d) == (3,) + @test innersize(A_similar_1d) == innersize(A) @test typeof(A_similar) == typeof(A) end From f280a0e79580d372a742cc5b7747c9f6c4a6b60e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:03:50 +0200 Subject: [PATCH 25/45] Make fused the interface function of AbstractArrayOfSimilarArrays Third-party subtypes of AbstractArrayOfSimilarArrays previously fell through to UnknownSplitMode, breaking flatview, stacked, innersize, etc. Now subtypes only need to implement fused (plus the standard array API), all split-mode operations are derived from it at the abstract type level. Also tighten the element type bound of the abstract type to AbstractArray{T,M}, which resolves the resulting method ambiguities. Co-Authored-By: Claude Fable 5 --- src/array_of_similar_arrays.jl | 91 +++++++++++++++++++-------------- test/array_of_similar_arrays.jl | 46 +++++++++++++++++ test/testdefs.jl | 2 + 3 files changed, 100 insertions(+), 39 deletions(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index e1fa6dc..4aa4234 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -1,16 +1,21 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). """ - AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T}} <: AbstractSlices{ET,N} + AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T,M}} <: AbstractSlices{ET,N} An array that contains arrays that have the same size/axes. The array is internally stored in flattened form as some kind of array of dimension -`M + N`. The flattened form can be accessed via `flatview(A)`. +`M + N`, in memory order. The flattened form can be accessed via +`flatview(A)`. -Subtypes must implement (in addition to typical array operations): +Subtypes must implement (in addition to typical array operations) - flatview(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N} - getslicemap(A::SomeArrayOfSimilarArrays) + ArraysOfArrays.fused(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N} + +which must return the underlying flat array. All split-mode operations +([`getsplitmode`](@ref), [`stacked`](@ref), [`flatview`](@ref), +[`innersize`](@ref), [`getslicemap`](@ref), etc.) are then provided +automatically. The following type aliases are defined: @@ -18,10 +23,10 @@ The following type aliases are defined: * `AbstractArrayOfSimilarVectors{T,N,ET} = AbstractArrayOfSimilarArrays{T,1,N,ET}` * `AbstractVectorOfSimilarVectors{T,ET} = AbstractArrayOfSimilarArrays{T,1,1,ET}` """ -abstract type AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T}} <: AbstractSlices{ET,N} end +abstract type AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T,M}} <: AbstractSlices{ET,N} end export AbstractArrayOfSimilarArrays -const AbstractVectorOfSimilarArrays{T,M,ET<:AbstractArray{T}} = AbstractArrayOfSimilarArrays{T,M,1,ET} +const AbstractVectorOfSimilarArrays{T,M,ET<:AbstractArray{T,M}} = AbstractArrayOfSimilarArrays{T,M,1,ET} export AbstractVectorOfSimilarArrays const AbstractArrayOfSimilarVectors{T,N,ET<:AbstractVector{T}} = AbstractArrayOfSimilarArrays{T,1,N,ET} @@ -57,6 +62,45 @@ getouterdims(obj::Tuple, ::SplitSlices{M,N}) where {M,N} = back_tuple(obj, Val(N @inline splitup(A::AbstractArray{T}, ::SplitSlices{M,N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(A) +# Subtypes of AbstractArrayOfSimilarArrays only need to implement `fused` +# (plus the typical array operations), all split-mode operations are derived +# from it: + +function fused(A::AbstractArrayOfSimilarArrays) + throw(ArgumentError("Subtypes of AbstractArrayOfSimilarArrays like $(nameof(typeof(A))) must implement ArraysOfArrays.fused")) +end + +@inline getsplitmode(::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} = SplitSlices{M,N}() +@inline unstackmode(A::AbstractArrayOfSimilarArrays) = getsplitmode(A) + +@inline stacked(A::AbstractArrayOfSimilarArrays) = fused(A) + +@inline Base.parent(A::AbstractArrayOfSimilarArrays) = fused(A) + +@inline vecflattened(A::AbstractArrayOfSimilarArrays) = vec(fused(A)) +# Disambiguation: +@inline vecflattened(A::AbstractVectorOfSimilarArrays) = vec(fused(A)) +@inline vecflattened(A::AbstractArrayOfSimilarVectors) = vec(fused(A)) +@inline vecflattened(A::AbstractVectorOfSimilarVectors) = vec(fused(A)) + +function getslicemap(::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} + return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) +end + +# `stack` must return an independent array, unlike `stacked`. Julia does not +# dispatch on keyword arguments, so dispatch on the value of `dims` instead: +Base.stack(A::AbstractArrayOfSimilarArrays; dims::Union{Integer,Colon} = :) = _stack_impl(A, dims) + +# Fast path for the default layout: a single bulk copy instead of the +# element-by-element copy that generic `stack` would do: +_stack_impl(A::AbstractArrayOfSimilarArrays, ::Colon) = copy(fused(A)) +_stack_impl(A::AbstractArrayOfSimilarArrays, dims::Integer) = stack(collect(A); dims) + +@inline Base.:(==)(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = (stacked(A) == stacked(B)) +@inline Base.isequal(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = isequal(stacked(A), stacked(B)) +@inline Base.isapprox(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}; kwargs...) where {M,N} = isapprox(stacked(A), stacked(B); kwargs...) + + """ ArrayOfSimilarArrays{T,M,N,P,ET} <: AbstractArrayOfSimilarArrays{T,M,N,ET} @@ -122,28 +166,7 @@ Base.convert(::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U Base.convert(::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) -@inline Base.:(==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} = (stacked(A) == stacked(B)) -@inline Base.isequal(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U} = isequal(stacked(A), stacked(B)) -@inline Base.isapprox(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{U,M,N}; kwargs...) where {T,M,N,U} = isapprox(stacked(A), stacked(B); kwargs...) - -@inline getsplitmode(A::ArrayOfSimilarArrays) = unstackmode(A) -@inline unstackmode(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = SplitSlices{M,N}() - -@inline stacked(A::ArrayOfSimilarArrays) = A.data -@inline fused(A::ArrayOfSimilarArrays) = stacked(A) - -@inline vecflattened(A::ArrayOfSimilarArrays) = vec(fused(A)) - -@inline Base.parent(A::ArrayOfSimilarArrays) = fused(A) - -# `stack` must return an independent array, unlike `stacked`. Julia does not -# dispatch on keyword arguments, so dispatch on the value of `dims` instead: -Base.stack(A::ArrayOfSimilarArrays; dims::Union{Integer,Colon} = :) = _stack_impl(A, dims) - -# Fast path for the default layout: a single bulk copy instead of the -# element-by-element copy that generic `stack` would do: -_stack_impl(A::ArrayOfSimilarArrays, ::Colon) = copy(fused(A)) -_stack_impl(A::ArrayOfSimilarArrays, dims::Integer) = stack(collect(A); dims) +@inline fused(A::ArrayOfSimilarArrays) = A.data function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} new_ET = Base.promote_op(similar, ET) @@ -151,16 +174,6 @@ function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} end -function getslicemap(::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} - return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) -end - - -@inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} - front_tuple(size(A.data), Val{M}()) -end - - Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index f02c42f..2b09c9c 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -10,6 +10,22 @@ using ChainRulesTestUtils using Statistics using StatsBase: cov2cor +include("testdefs.jl") + +# Minimal third-party subtype of AbstractArrayOfSimilarArrays, only +# implements the standard array API and `fused`: +if !isdefined(Main, :TestSimilarVectors) + struct TestSimilarVectors{T,ET} <: AbstractArrayOfSimilarArrays{T,1,1,ET} + data::Matrix{T} + + TestSimilarVectors(data::Matrix{T}) where {T} = + new{T,typeof(view(data, :, firstindex(data, 2)))}(data) + end + Base.size(A::TestSimilarVectors) = (size(A.data, 2),) + Base.getindex(A::TestSimilarVectors, i::Int) = view(A.data, :, i) + ArraysOfArrays.fused(A::TestSimilarVectors) = A.data +end + @testset "array_of_similar_arrays" begin function rand_flat_array(Val_N::Val{N}) where {N} sz_max = (2,3,2,4,5) @@ -149,6 +165,36 @@ using StatsBase: cov2cor end end + @testset "split mode API" begin + A_flat = rand_flat_array(Val(5)) + A = ArrayOfSimilarArrays{Float64,2,3}(A_flat) + test_api(A, Array(A), A_flat) + + V_flat = rand_flat_array(Val(2)) + V = VectorOfSimilarVectors(V_flat) + test_api(V, Array(V), V_flat) + end + + @testset "custom subtypes" begin + data = rand(3, 5) + B = TestSimilarVectors(data) + + @test @inferred(getsplitmode(B)) === SplitSlices{1,1}() + @test @inferred(unstackmode(B)) === SplitSlices{1,1}() + @test @inferred(fused(B)) === data + @test @inferred(stacked(B)) === data + @test @inferred(flatview(B)) === data + @test @inferred(parent(B)) === data + @test @inferred(vecflattened(B)) == vec(data) + @test @inferred(innersize(B)) == (3,) + @test @inferred(getslicemap(B)) == (:, 1) + @test @inferred(stack(B)) == data + @test stack(B) !== data + @test B == VectorOfSimilarVectors(data) + @test isapprox(B, VectorOfSimilarVectors(data .+ 1e-14)) + @test splitup(fused(B), getsplitmode(B)) == B + end + @testset "add remove" begin EA_ref1 = rand_flat_array(Val(3)) EA_ref2 = rand_flat_array(Val(3)) diff --git a/test/testdefs.jl b/test/testdefs.jl index 099a4d6..561f168 100644 --- a/test/testdefs.jl +++ b/test/testdefs.jl @@ -2,6 +2,8 @@ using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent +using ArraysOfArrays: getinnerdims, getouterdims + if !isdefined(Main, :test_api) maptest_f(x::Number) = x^2 maptest_f(x::AbstractArray{<:Number}) = sum(x)^2 From 7811c262283b6c3ddc0a8ef10e54134812639454 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:08:10 +0200 Subject: [PATCH 26/45] Restrict split-mode API to AbstractArray getsplitmode, unstackmode, fused, splitup, innermap and deepmap had generic Any fallbacks that gave non-array objects an UnknownSplitMode and confusing split-mode errors. Non-arrays now get a plain MethodError. Co-Authored-By: Claude Fable 5 --- src/functions.jl | 34 ++++++++++++++-------------------- 1 file changed, 14 insertions(+), 20 deletions(-) diff --git a/src/functions.jl b/src/functions.jl index b27812c..015b011 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -95,8 +95,6 @@ the same type as `A` if at all possible, except if `getsplitmode(A)` is an function getsplitmode end export getsplitmode -@inline getsplitmode(::T) where T = UnknownSplitMode{T}() - @inline getsplitmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() @inline getsplitmode(A::AbstractArray{<:AbstractArray}) = UnknownSplitMode{typeof(A)}() @@ -114,9 +112,9 @@ See also [`fused`](@ref) and [`getsplitmode`](@ref). function splitup end export splitup -@inline splitup(obj::Any, ::NonSplitMode) = obj +@inline splitup(A::AbstractArray, ::NonSplitMode) = A -function splitup(::Any, ::UnknownSplitMode) +function splitup(::AbstractArray, ::UnknownSplitMode) throw(ArgumentError("splitup cannot be used with UnknownSplitMode")) end @@ -142,8 +140,6 @@ equivalent to [`flatview(A)`](@ref). function fused end export fused -@inline fused(obj) = _fused_impl(obj, getsplitmode(obj)) - @inline fused(A::AbstractArray) = A @inline fused(A::AbstractArray{<:AbstractArray}) = _fused_impl(A, getsplitmode(A)) @@ -267,10 +263,10 @@ type and underlying memory layout than `A`. function unstackmode end export unstackmode -@inline unstackmode(::T) where T = UnknownSplitMode{T}() - @inline unstackmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() +@inline unstackmode(A::AbstractArray{<:AbstractArray}) = UnknownSplitMode{typeof(A)}() + function unstackmode(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} innersize(A) # Ensure element arrays have equal size return SplitSlices{M,N}() @@ -367,16 +363,15 @@ of arrays, otherwise equivalent to `Base.map`. function innermap end export innermap -innermap(f, obj) = _generic_innermap_impl(f, obj) innermap(f, A::AbstractArray) = map(f, A) innermap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(map, f), A) innermap(f, A::AbstractSlices{<:AbstractArray}) = _generic_innermap_impl(f, A) -function _generic_innermap_impl(f, obj) - joined_obj = fused(obj) - mapped_joined_obj = map(f, joined_obj) - mapped_obj = splitup(mapped_joined_obj, getsplitmode(obj)) - return mapped_obj +function _generic_innermap_impl(f, A::AbstractArray) + joined_A = fused(A) + mapped_joined_A = map(f, joined_A) + mapped_A = splitup(mapped_joined_A, getsplitmode(A)) + return mapped_A end @@ -443,16 +438,15 @@ a nested array, `deepmap` behaves identically to `Base.map`. function deepmap end export deepmap -deepmap(f, obj) = _generic_deepmap_impl(f, obj) deepmap(f, A::AbstractArray) = map(f, A) deepmap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(deepmap, f), A) deepmap(f, A::AbstractSlices{<:AbstractArray}) = _generic_deepmap_impl(f, A) -function _generic_deepmap_impl(f, obj) - joined_obj = fused(obj) - mapped_joined_obj = deepmap(f, joined_obj) - mapped_obj = splitup(mapped_joined_obj, getsplitmode(obj)) - return mapped_obj +function _generic_deepmap_impl(f, A::AbstractArray) + joined_A = fused(A) + mapped_joined_A = deepmap(f, joined_A) + mapped_A = splitup(mapped_joined_A, getsplitmode(A)) + return mapped_A end From dcaafedcf9b082bce4e471d415714f27767a44c0 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:11:16 +0200 Subject: [PATCH 27/45] Copy shape information in getsplitmode of VectorOfArrays The split mode shared elem_ptr and kernel_size with the vector of arrays it came from, so resizing the vector silently mutated the mode. getsplitmode now makes a defensive copy (package extensions may override this for vector types that cannot be resized). SplitParts gains structural == and hash, since its instances no longer share field identity. Co-Authored-By: Claude Fable 5 --- src/vector_of_arrays.jl | 19 ++++++++++++++++++- test/vector_of_arrays.jl | 6 ++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index e3573fb..4700682 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -238,6 +238,12 @@ SplitParts( `elem_ptr` and `kernel_size` equal the equivalent properties of `VectorOfArrays`. +[`getsplitmode(A::VectorOfArrays)`](@ref) copies the shape information of +`A` (an O(length) operation), so the resulting mode is not affected if `A` +is resized afterwards. A `VectorOfArrays` created via +[`splitup`](@ref), on the other hand, shares the vectors of the mode it was +created from, like the `VectorOfArrays` inner constructor. + See also [`AbstractPartMode`](@ref). """ struct SplitParts{ @@ -264,7 +270,18 @@ export SplitParts is_memordered_splitmode(::SplitParts) = true -@inline getsplitmode(A::VectorOfArrays) = SplitParts(A.elem_ptr, A.kernel_size) +Base.:(==)(a::SplitParts, b::SplitParts) = + a.elem_ptr == b.elem_ptr && a.kernel_size == b.kernel_size + +Base.hash(x::SplitParts, h::UInt) = + hash(x.kernel_size, hash(x.elem_ptr, hash(:SplitParts, h))) + +# Defensive copy: the split mode must not be affected if the vector of +# arrays is resized later on. Package extensions may specialize this for +# vector types that cannot be resized: +_shapeinfo_copy(x::AbstractVector) = copy(x) + +getsplitmode(A::VectorOfArrays) = SplitParts(_shapeinfo_copy(A.elem_ptr), _shapeinfo_copy(A.kernel_size)) @inline fused(A::VectorOfArrays) = A.data diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 3331e11..4d8eab4 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -154,6 +154,12 @@ include("testdefs.jl") test_api(B, Array(B), B.data) end + # getsplitmode must not be affected by later resizing: + B_grow = VectorOfArrays([[1, 2], [3, 4, 5]]) + sm_grow = getsplitmode(B_grow) + push!(B_grow, [6]) + @test splitup(collect(1:5), sm_grow) == [[1, 2], [3, 4, 5]] + # Uniform element size, so stackable: @test @inferred(stacked(Bu)) == stack(Array(Bu)) @test @inferred(splitup(stacked(Bu), unstackmode(Bu))) == Bu From 480713ebb002d7ce3c4ef3cd50e74ff38dab51f7 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:14:29 +0200 Subject: [PATCH 28/45] Support dims keyword in reductions over VectorOfSimilarArrays sum, mean, var and std specializations had no keyword arguments, so calls like mean(X; dims = 1) threw a MethodError instead of reaching the generic implementations. Dispatch on the value of dims instead, keeping the fast flat-data path for dims = Colon. Also add fill!(::ArrayOfSimilarArrays, x::AbstractArray), which the generic reduction machinery requires (and which is useful on its own). Co-Authored-By: Claude Fable 5 --- src/array_of_similar_arrays.jl | 36 +++++++++++++++++++++++++++++---- test/array_of_similar_arrays.jl | 6 ++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 4aa4234..06d46dc 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -184,6 +184,12 @@ Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Var Base.@propagate_inbounds Base.setindex!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = setindex!(A.data, x, _ncolons(Val{M}())..., idxs...) +function Base.fill!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}) where {T,M,N,U} + size(x) == innersize(A) || throw(DimensionMismatch("Can't fill ArrayOfSimilarArrays with an array of different size than its elements")) + A.data .= reshape(x, size(x)..., ntuple(_ -> 1, Val(N))...) + return A +end + Base.@propagate_inbounds function Base.unsafe_view(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} dataview = view(A.data, _ncolons(Val{M}())..., idxs...) L = length(size(dataview)) @@ -344,17 +350,39 @@ Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVec Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(stacked(A)) -Base.sum(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = +# Fast paths over the flat data for the full (dims = :) reductions. Julia +# does not dispatch on keyword arguments, so dispatch on the value of `dims` +# instead and forward non-Colon dims to the generic implementations: + +Base.sum(X::AbstractVectorOfSimilarArrays; dims::Union{Colon,Integer} = :) = _aosa_sum(X, dims) + +_aosa_sum(X::AbstractVectorOfSimilarArrays{T,M}, ::Colon) where {T,M} = sum(flatview(X); dims = M + 1)[_ncolons(Val{M}())...] +_aosa_sum(X::AbstractVectorOfSimilarArrays, dims::Integer) = + Base.@invoke sum(X::AbstractArray; dims = dims) -Statistics.mean(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = +Statistics.mean(X::AbstractVectorOfSimilarArrays; dims::Union{Colon,Integer} = :) = _aosa_mean(X, dims) + +_aosa_mean(X::AbstractVectorOfSimilarArrays{T,M}, ::Colon) where {T,M} = mean(flatview(X); dims = M + 1)[_ncolons(Val{M}())...] +_aosa_mean(X::AbstractVectorOfSimilarArrays, dims::Integer) = + Base.@invoke mean(X::AbstractArray; dims = dims) + +Statistics.var(X::AbstractVectorOfSimilarArrays; corrected::Bool = true, mean = nothing, dims::Union{Colon,Integer} = :) = + _aosa_var(X, corrected, mean, dims) -Statistics.var(X::AbstractVectorOfSimilarArrays{T,M}; corrected::Bool = true, mean = nothing) where {T,M} = +_aosa_var(X::AbstractVectorOfSimilarArrays{T,M}, corrected::Bool, mean, ::Colon) where {T,M} = var(flatview(X); dims = M + 1, corrected = corrected, mean = mean)[_ncolons(Val{M}())...] +_aosa_var(X::AbstractVectorOfSimilarArrays, corrected::Bool, mean, dims::Integer) = + Base.@invoke var(X::AbstractArray; corrected = corrected, mean = mean, dims = dims) + +Statistics.std(X::AbstractVectorOfSimilarArrays; corrected::Bool = true, mean = nothing, dims::Union{Colon,Integer} = :) = + _aosa_std(X, corrected, mean, dims) -Statistics.std(X::AbstractVectorOfSimilarArrays{T,M}; corrected::Bool = true, mean = nothing) where {T,M} = +_aosa_std(X::AbstractVectorOfSimilarArrays{T,M}, corrected::Bool, mean, ::Colon) where {T,M} = std(flatview(X); dims = M + 1, corrected = corrected, mean = mean)[_ncolons(Val{M}())...] +_aosa_std(X::AbstractVectorOfSimilarArrays, corrected::Bool, mean, dims::Integer) = + Base.@invoke std(X::AbstractArray; corrected = corrected, mean = mean, dims = dims) Statistics.cov(X::AbstractVectorOfSimilarVectors; corrected::Bool = true) = cov(flatview(X); dims = 2, corrected = corrected) diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 2b09c9c..f39e00c 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -292,6 +292,12 @@ end VA = [rand(2,3,3) for i in 1:10] VA_aosa = ArrayOfSimilarArrays(VA) + # Non-Colon dims must forward to the generic implementations: + @test sum(VV_aosa; dims = 1) == sum(collect(VV_aosa); dims = 1) + @test mean(VV_aosa; dims = 1) == mean(collect(VV_aosa); dims = 1) + @test @inferred(sum(VV_aosa)) ≈ sum(collect(VV_aosa)) + @test @inferred(mean(VV_aosa)) ≈ mean(collect(VV_aosa)) + array_cmp(A, B) = (A ≈ B) && (size(A) == size(B)) function test_statistics_op(op::Function) From 8b9a8580fec09f39571060a6b068cb8f8606f625 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:15:03 +0200 Subject: [PATCH 29/45] Document the flattening function family Add a short guide on fused vs flatview vs stacked vs vecflattened and their memory-sharing behavior, and note the SVector variant of sliced in its docstring. Co-Authored-By: Claude Fable 5 --- docs/src/index.md | 9 +++++++++ src/functions.jl | 3 +++ 2 files changed, 12 insertions(+) diff --git a/docs/src/index.md b/docs/src/index.md index 1329754..3a12764 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -13,6 +13,15 @@ This package also defines and exports the following new functions applicable to * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. +## Which flattening function do I want? + +* [`fused(A)`](@ref): the original underlying array, `splitup(fused(A), getsplitmode(A)) == A`. +* [`flatview(A)`](@ref): the underlying storage, requires a memory-ordered layout. +* [`stacked(A)`](@ref): elements joined along new trailing dimensions, like `Base.stack`. +* [`vecflattened(A)`](@ref): elements concatenated into a single vector, like `reduce(vcat, A)`. + +All four are zero-copy where possible and so may return arrays that share memory with `A`. In contrast, `Base.stack(A)` and `reduce(vcat, A)` always return independent arrays. + ## [ArrayOfSimilarArrays](@id section_ArrayOfSimilarArrays) diff --git a/src/functions.jl b/src/functions.jl index 015b011..7b15a85 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -310,6 +310,9 @@ end Return a sliced view of `A`, using the columns or the first `M` dimensions as inner dimensions. + +With StaticArrays loaded, `sliced(A, SVector{S})` returns a reinterpreted +array with `SVector{S}` elements instead of an array of views. """ function sliced end export sliced From 479339c0f601dbfd0e8a7dbfd0584e8848852a08 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:17:00 +0200 Subject: [PATCH 30/45] Add Adapt support for SplitParts SplitParts carries elem_ptr and kernel_size vectors, so split modes must be adaptable to other array storage (e.g. GPU) like the vectors of arrays they describe. Co-Authored-By: Claude Fable 5 --- ext/ArraysOfArraysAdaptExt.jl | 10 +++++++++- test/vector_of_arrays.jl | 4 ++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/ext/ArraysOfArraysAdaptExt.jl b/ext/ArraysOfArraysAdaptExt.jl index 6bd2e65..69f37c0 100644 --- a/ext/ArraysOfArraysAdaptExt.jl +++ b/ext/ArraysOfArraysAdaptExt.jl @@ -5,7 +5,7 @@ module ArraysOfArraysAdaptExt import Adapt using Adapt: adapt -using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays +using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays, SplitParts using ArraysOfArrays: no_consistency_checks @@ -25,4 +25,12 @@ function Adapt.adapt_structure(to, A::VectorOfArrays) end +function Adapt.adapt_structure(to, smode::SplitParts) + SplitParts( + adapt(to, smode.elem_ptr), + adapt(to, smode.kernel_size) + ) +end + + end # module ArraysOfArraysAdaptExt diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 4d8eab4..42058ce 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -340,6 +340,10 @@ include("testdefs.jl") A3 = VectorOfArrays(ref_AoA3(Float32, 3)) @test @inferred(adapt(identity, A3)) == A3 @test typeof(adapt(identity, A3)) == typeof(A3) + + sm = getsplitmode(A3) + @test @inferred(adapt(identity, sm)) == sm + @test typeof(adapt(identity, sm)) == typeof(sm) end From 7b0ad0c245966a2dd34e1948f0187db272b8d2f6 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:22:20 +0200 Subject: [PATCH 31/45] Add ChainRules rrules for partitioned and vecflattened The pullbacks reconstruct nested tangents via splitup and give no tangent contribution to data not covered by partial partitions. Co-Authored-By: Claude Fable 5 --- ext/ArraysOfArraysChainRulesCoreExt.jl | 61 +++++++++++++++++++++++++- test/array_of_similar_arrays.jl | 9 ++++ test/vector_of_arrays.jl | 30 +++++++++++++ 3 files changed, 98 insertions(+), 2 deletions(-) diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index c46ef92..83a6863 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -5,9 +5,9 @@ module ArraysOfArraysChainRulesCoreExt using ChainRulesCore: ChainRulesCore, NoTangent, AbstractThunk, Thunk, unthunk, @thunk, @non_differentiable using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitup, fused, stacked, unstackmode, - flatview, innersize + flatview, innersize, vecflattened, partitioned using ArraysOfArrays: NonSplitMode, AbstractSplitMode -using ArraysOfArrays: AbstractArrayOfSimilarArrays, ArrayOfSimilarArrays +using ArraysOfArrays: AbstractArrayOfSimilarArrays, ArrayOfSimilarArrays, VectorOfArrays struct _MappedMaybeThunk{F, T} <: AbstractThunk @@ -59,6 +59,63 @@ _stacked_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitup, smode) +function ChainRulesCore.rrule(::typeof(vecflattened), A::AbstractArrayOfSimilarArrays) + smode = getsplitmode(A) + sz = size(fused(A)) + function aosa_vecflattened_pullback(ΔΩ) + Δdata = reshape(unthunk(ΔΩ), sz) + return NoTangent(), splitup(Δdata, smode) + end + return vecflattened(A), aosa_vecflattened_pullback +end + +function ChainRulesCore.rrule(::typeof(vecflattened), A::VectorOfArrays) + smode = getsplitmode(A) # makes a defensive copy, safe to close over + datalen = length(A.data) + covered_from = first(A.elem_ptr) + covered_len = last(A.elem_ptr) - covered_from + function voa_vecflattened_pullback(ΔΩ) + Δ = unthunk(ΔΩ) + Δdata = if covered_len == datalen + Δ + else + padded = fill!(similar(Δ, datalen), zero(eltype(Δ))) + copyto!(padded, covered_from, Δ, firstindex(Δ), covered_len) + padded + end + return NoTangent(), splitup(Δdata, smode) + end + return vecflattened(A), voa_vecflattened_pullback +end + + +function ChainRulesCore.rrule(::typeof(partitioned), A::AbstractVector, lengths::AbstractVector{<:Integer}) + return partitioned(A, lengths), _partitioned_pullback_for(A) +end + +function ChainRulesCore.rrule(::typeof(partitioned), A::AbstractVector, shapes::AbstractVector{<:Dims}) + return partitioned(A, shapes), _partitioned_pullback_for(A) +end + +function _partitioned_pullback_for(A::AbstractVector) + datalen = length(A) + function partitioned_pullback(ΔΩ) + Δflat = vecflattened(unthunk(ΔΩ)) + ΔA = if length(Δflat) == datalen + Δflat + else + # Partial partitions leave uncovered data, which gets no tangent + # contribution: + padded = fill!(similar(Δflat, datalen), zero(eltype(Δflat))) + copyto!(padded, firstindex(padded), Δflat, firstindex(Δflat), length(Δflat)) + padded + end + return NoTangent(), ΔA, NoTangent() + end + return partitioned_pullback +end + + function _aosa_ctor_fromflat_pullback(ΔΩ) NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) end diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index f39e00c..fd229e0 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -173,6 +173,15 @@ end V_flat = rand_flat_array(Val(2)) V = VectorOfSimilarVectors(V_flat) test_api(V, Array(V), V_flat) + + # vecflattened rrule: + A_rr = ArrayOfSimilarArrays{Float64,1,1}(rand(3, 4)) + y, pb = rrule(vecflattened, A_rr) + @test y == vec(A_rr.data) + t = pb(collect(1.0:12.0)) + @test t[1] == NoTangent() + @test t[2] isa ArrayOfSimilarArrays + @test fused(t[2]) == reshape(1.0:12.0, 3, 4) end @testset "custom subtypes" begin diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 42058ce..4d3a98d 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -208,6 +208,36 @@ include("testdefs.jl") end + @testset "rrules" begin + x = collect(Float64, 1:10) + + for lengths in ([2, 3, 5], [2, 3]) + Y, pb = rrule(partitioned, x, lengths) + @test Y == partitioned(x, lengths) + ΔY = [fill(1.0, l) for l in lengths] + ΔA = pb(ΔY)[2] + want = zero(x) + want[1:sum(lengths)] .= 1.0 + @test pb(ΔY)[1] == NoTangent() && pb(ΔY)[3] == NoTangent() + @test ΔA == want + end + + Y2, pb2 = rrule(partitioned, x, [(1, 2), (2, 2)]) + @test Y2 == partitioned(x, [(1, 2), (2, 2)]) + ΔA2 = pb2([fill(1.0, 1, 2), fill(1.0, 2, 2)])[2] + @test ΔA2 == [ones(6); zeros(4)] + + p = partitioned(x, [2, 3]) + y, pb3 = rrule(vecflattened, p) + @test y == vecflattened(p) + t = pb3(collect(1.0:5.0)) + @test t[1] == NoTangent() + @test t[2] isa VectorOfArrays + @test t[2] == [[1.0, 2.0], [3.0, 4.0, 5.0]] + @test fused(t[2])[6:10] == zeros(5) + end + + @testset "indexing" begin V1 = @inferred(VectorOfArrays(ref_AoA3(Float32, 3))) V2 = @inferred(VectorOfArrays(ref_AoA3(Float32, 3))) From 723de1100e223059db9195d07a86c5f1db21322a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:27:10 +0200 Subject: [PATCH 32/45] Fast single-pass vcat and reduce(vcat) for vectors of arrays - vcat and reduce(vcat, ...) of VectorOfSimilarArrays now concatenate the underlying data along its last dimension with a single allocation and return a VectorOfSimilarArrays (previously a generic Vector of views) - reduce(vcat, ...) of VectorOfArrays uses a single-pass implementation instead of pairwise reduction - vcat and append! of VectorOfArrays now only use the data covered by the element arrays, both silently corrupted the result for inputs with unused data regions (e.g. from partial partitions). append! to a target with unused trailing data now throws. - vcat(V::VectorOfArrays) returned V itself, it now returns an independent copy like Base.vcat This (mostly) upstreams LegendDataTypes.fast_flatten: it now amounts to reduce(vcat, xs) for vectors of ArraysOfArrays types, and to flatview(reduce(vcat, map(x -> sliced(x, Val(ndims(x) - 1)), xs))) for generic arrays concatenated along their last dimension. Co-Authored-By: Claude Fable 5 --- src/array_of_similar_arrays.jl | 14 +++++++++ src/util.jl | 26 ++++++++++++++++ src/vector_of_arrays.jl | 56 ++++++++++++++++++++++++++------- test/vector_of_arrays.jl | 57 +++++++++++++++++++++++++++++++++- 4 files changed, 141 insertions(+), 12 deletions(-) diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 06d46dc..349d655 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -301,6 +301,20 @@ end # popfirst!(V::ArrayOfSimilarArrays) = ... +# Concatenating vectors of similar arrays concatenates their underlying +# data along its last dimension, with a single allocation: + +function _vcat_vosas(Vs) + isempty(Vs) && throw(ArgumentError("reducing over an empty collection is not allowed")) + new_data = _cat_lastdim(map(fused, Vs)) + return VectorOfSimilarArrays(new_data) +end + +Base.vcat(V1::AbstractVectorOfSimilarArrays, Vs::AbstractVectorOfSimilarArrays...) = _vcat_vosas((V1, Vs...)) + +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:AbstractVectorOfSimilarArrays}) = _vcat_vosas(Vs) + + function _empty_data_size(A::VectorOfSimilarArrays{T,M}) where {T,M} size_inner, size_outer = split_tuple(size(A.data), Val{M}()) empty_size_outer = map(x -> zero(x), size_outer) diff --git a/src/util.jl b/src/util.jl index b3034d5..adde6d6 100644 --- a/src/util.jl +++ b/src/util.jl @@ -41,3 +41,29 @@ function _require_ndims(::Val{N1}, ::Val{N2}) where {N1,N2} end Base.@pure _val_value(::Val{x}) where x = x + + +# Concatenate arrays of equal ndims and equal size except in their last +# dimension along their last dimension, with a single allocation: +_cat_lastdim(datas) = _cat_lastdim_impl(datas, Val(ndims(first(datas)))) + +function _cat_lastdim_impl(datas, ::Val{N}) where {N} + data1 = first(datas) + inner_sz = Base.front(size(data1)) + foreach(datas) do d + ndims(d) == N && Base.front(size(d)) == inner_sz || throw(DimensionMismatch("Can't concatenate arrays with different sizes in their non-last dimensions along their last dimension")) + end + + T = mapreduce(eltype, promote_type, datas) + n_lastdim = sum(d -> size(d, N), datas) + result = similar(data1, T, (inner_sz..., n_lastdim)) + + colons = ntuple(_ -> :, Val(N - 1)) + offset = firstindex(result, N) + for d in datas + n = size(d, N) + result[colons..., offset:(offset + n - 1)] = d + offset += n + end + return result +end diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 4700682..1b053f8 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -452,7 +452,11 @@ end function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U} if !isempty(B) - append!(A.data, B.data) + # append_elemptr! places the elements of B directly after the last + # element of A, so A must not contain unused trailing data and only + # the data covered by the elements of B may be appended: + last(A.elem_ptr) == lastindex(A.data) + 1 || throw(ArgumentError("Cannot append to a VectorOfArrays that has unused trailing data")) + append!(A.data, vecflattened(B)) append_elemptr!(A.elem_ptr, B.elem_ptr) append!(A.kernel_size, B.kernel_size) end @@ -482,22 +486,52 @@ function Base.append!(A::VectorOfArrays{T,N}, B::AbstractVector{<:AbstractArray{ end -Base.vcat(V::VectorOfArrays) = V +# Concatenating vectors of arrays concatenates the data covered by their +# elements, in a single pass: -function Base.vcat(Vs::(VectorOfArrays{U,N} where U)...) where {N} - data = vcat(map(x -> x.data, Vs)...) +function _vcat_voas(Vs) + isempty(Vs) && throw(ArgumentError("reducing over an empty collection is not allowed")) + V1 = first(Vs) - elem_ptr = similar(Vs[1].elem_ptr, 1) - elem_ptr[1] = firstindex(data) - @inbounds for i in eachindex(Vs) - append_elemptr!(elem_ptr, Vs[i].elem_ptr) - end + T = mapreduce(V -> eltype(V.data), promote_type, Vs) + n_parts = sum(length, Vs) + n_data = sum(V -> Int(last(V.elem_ptr) - first(V.elem_ptr)), Vs) - kernel_size = vcat(map(x -> x.kernel_size, Vs)...) + data = similar(V1.data, T, n_data) + elem_ptr = similar(V1.elem_ptr, n_parts + 1) + kernel_size = similar(V1.kernel_size, n_parts) - VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks) + i_ptr = firstindex(elem_ptr) + elem_ptr[i_ptr] = firstindex(data) + i_ksz = firstindex(kernel_size) + i_data = firstindex(data) + + for V in Vs + ep = V.elem_ptr + covered_len = last(ep) - first(ep) + copyto!(data, i_data, V.data, first(ep), covered_len) + i_data += covered_len + + @inbounds for j in firstindex(ep):(lastindex(ep) - 1) + elem_ptr[i_ptr + 1] = elem_ptr[i_ptr] + (ep[j + 1] - ep[j]) + i_ptr += 1 + end + + ksz = V.kernel_size + copyto!(kernel_size, i_ksz, ksz, firstindex(ksz), length(ksz)) + i_ksz += length(ksz) + end + + return VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks) end +Base.vcat(V1::VectorOfArrays{<:Any,N}, Vs::VectorOfArrays{<:Any,N}...) where {N} = _vcat_voas((V1, Vs...)) + +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays}) = _vcat_voas(Vs) +# Disambiguation: +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays{<:Any,1}}) = _vcat_voas(Vs) +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays{<:Any,2}}) = _vcat_voas(Vs) + function Base.copy(V::VectorOfArrays) VectorOfArrays(copy(V.data), copy(V.elem_ptr), copy(V.kernel_size), no_consistency_checks) diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 4d3a98d..e247beb 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -85,7 +85,10 @@ include("testdefs.jl") B1 = VectorOfArrays(A1); B2 = VectorOfArrays(A2); B3 = VectorOfArrays(A3); B4 = VectorOfArrays(A4); - @test @inferred(vcat(B1)) === B1 + # vcat must return an independent array: + @test @inferred(vcat(B1)) == B1 + @test vcat(B1) !== B1 + @test vcat(B1).data !== B1.data full_consistency_checks(vcat(B1)) @test @inferred(vcat(B1, B2)) isa VectorOfArrays @@ -107,6 +110,58 @@ include("testdefs.jl") B1_copy = @inferred(copy(B1)); B3_copy = @inferred(copy(B3)) append!(B1_copy, B3_copy) @test B1_copy.data == vcat(B1.data, B3.data) + + # reduce(vcat, ...) uses a single-pass implementation: + @test @inferred(reduce(vcat, [B1, B2, B3])) isa VectorOfArrays + @test reduce(vcat, [B1, B2, B3]) == vcat(A1, A2, A3) + @test reduce(vcat, [B1, B2, B3]) == vcat(B1, B2, B3) + full_consistency_checks(reduce(vcat, [B1, B2, B3])) + + # vcat and append! of vectors of arrays with unused data regions + # must only use the data covered by their elements: + p_partial = partitioned(collect(1:10), [2, 3]) + q = VectorOfArrays([[6], [7, 8]]) + @test vcat(p_partial, q) == vcat(collect(p_partial), collect(q)) + @test reduce(vcat, [p_partial, q]) == vcat(collect(p_partial), collect(q)) + full_consistency_checks(vcat(p_partial, q)) + + q_copy = copy(q) + append!(q_copy, p_partial) + @test q_copy == vcat(collect(q), collect(p_partial)) + full_consistency_checks(q_copy) + + # Appending to a VectorOfArrays with unused trailing data would + # corrupt it: + @test_throws ArgumentError append!(partitioned(collect(1:10), [2, 3]), q) + end + + + @testset "vcat of VectorOfSimilarArrays" begin + VsoA = [VectorOfSimilarArrays(rand(Float32, 2, 3, n)) for n in (2, 0, 4)] + Vs_ref = vcat(map(collect, VsoA)...) + + @test @inferred(vcat(VsoA[1])) isa VectorOfSimilarArrays + @test vcat(VsoA[1]) == VsoA[1] + @test vcat(VsoA[1]).data !== VsoA[1].data + + @test @inferred(vcat(VsoA...)) isa VectorOfSimilarArrays + @test vcat(VsoA...) == Vs_ref + @test @inferred(reduce(vcat, VsoA)) isa VectorOfSimilarArrays + @test reduce(vcat, VsoA) == Vs_ref + + # Element type promotion like Base.vcat: + V64 = VectorOfSimilarArrays(rand(Float64, 2, 3, 2)) + @test eltype(eltype(vcat(VsoA[1], V64))) == Float64 + @test vcat(VsoA[1], V64) == vcat(collect(VsoA[1]), collect(V64)) + + # Inner size mismatch: + V_bad = VectorOfSimilarArrays(rand(Float32, 3, 2, 2)) + @test_throws DimensionMismatch vcat(VsoA[1], V_bad) + + # Vectors of similar vectors: + VsoV = [VectorOfSimilarVectors(rand(3, n)) for n in (2, 4)] + @test @inferred(reduce(vcat, VsoV)) isa VectorOfSimilarVectors + @test reduce(vcat, VsoV) == vcat(map(collect, VsoV)...) end From c87c50dd174d920c355e35ce627621ac9cc07620 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:31:57 +0200 Subject: [PATCH 33/45] Add FixedSizeArrays extension FixedSizeVectors cannot be resized, so a VectorOfArrays that uses them for elem_ptr and kernel_size has an immutable shape and getsplitmode can safely share the shape information instead of copying it. partitioned(A, lengths::FixedSizeVector) automatically results in fixed-size shape information via similar. Co-Authored-By: Claude Fable 5 --- Project.toml | 3 +++ ext/ArraysOfArraysFixedSizeArraysExt.jl | 14 ++++++++++ test/Project.toml | 1 + test/fixed_size_arrays.jl | 35 +++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 54 insertions(+) create mode 100644 ext/ArraysOfArraysFixedSizeArraysExt.jl create mode 100644 test/fixed_size_arrays.jl diff --git a/Project.toml b/Project.toml index cd3e82c..e2039d8 100644 --- a/Project.toml +++ b/Project.toml @@ -8,16 +8,19 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [extensions] ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" +ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" [compat] Adapt = "1, 2, 3, 4" ChainRulesCore = "1" +FixedSizeArrays = "1" StaticArraysCore = "1" Statistics = "1" julia = "1.10" diff --git a/ext/ArraysOfArraysFixedSizeArraysExt.jl b/ext/ArraysOfArraysFixedSizeArraysExt.jl new file mode 100644 index 0000000..a740024 --- /dev/null +++ b/ext/ArraysOfArraysFixedSizeArraysExt.jl @@ -0,0 +1,14 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysFixedSizeArraysExt + +using FixedSizeArrays: FixedSizeVector + +import ArraysOfArrays + +# FixedSizeVectors cannot be resized, so sharing them between a +# VectorOfArrays and a split mode is safe and no defensive copy is +# required: +ArraysOfArrays._shapeinfo_copy(x::FixedSizeVector) = x + +end # module ArraysOfArraysFixedSizeArraysExt diff --git a/test/Project.toml b/test/Project.toml index ea0ca2f..221c071 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -5,6 +5,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" +FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" diff --git a/test/fixed_size_arrays.jl b/test/fixed_size_arrays.jl new file mode 100644 index 0000000..41d611b --- /dev/null +++ b/test/fixed_size_arrays.jl @@ -0,0 +1,35 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using FixedSizeArrays: FixedSizeVector + +include("testdefs.jl") + +@testset "FixedSizeArrays extension" begin + A = VectorOfArrays([[1, 2], [3, 4, 5]]) + + # Fixed-size shape information cannot be resized, so getsplitmode can + # share it without a defensive copy: + B = VectorOfArrays(A.data, FixedSizeVector(A.elem_ptr), FixedSizeVector(A.kernel_size)) + sm = getsplitmode(B) + @test sm.elem_ptr === B.elem_ptr + @test sm.kernel_size === B.kernel_size + @test splitup(fused(B), sm) == B + @test typeof(splitup(fused(B), sm)) == typeof(B) + @test B == A + test_api(B, Array(B), B.data) + + # Resizable shape information still gets copied: + sm_A = getsplitmode(A) + @test sm_A.elem_ptr !== A.elem_ptr + @test sm_A.kernel_size !== A.kernel_size + + # partitioned with fixed-size lengths results in fixed-size shape + # information: + p = partitioned(collect(1:10), FixedSizeVector([2, 3, 5])) + @test p == [[1, 2], [3, 4, 5], [6, 7, 8, 9, 10]] + @test p.elem_ptr isa FixedSizeVector + @test getsplitmode(p).elem_ptr === p.elem_ptr +end diff --git a/test/runtests.jl b/test/runtests.jl index cabda5d..fc21dcc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,5 +10,6 @@ Test.@testset "Package ArraysOfArrays" begin include("vector_of_arrays.jl") include("broadcasting.jl") include("static_arrays.jl") + include("fixed_size_arrays.jl") include("test_docs.jl") end # testset From 8f5ac02c0f924ed7ca1abdfdeacba993906897de Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 02:32:09 +0200 Subject: [PATCH 34/45] Bump version to 0.7.0 The element-type and split-mode rework is breaking. Co-Authored-By: Claude Fable 5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index e2039d8..e11f190 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ArraysOfArrays" uuid = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" -version = "0.6.6" +version = "0.7.0" [deps] Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" From 4cc745fe5a6194ccea53763159382238f0a82938 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 08:27:52 +0200 Subject: [PATCH 35/45] Add Mooncake extension Mooncake differentiates the ArraysOfArrays types and operations without custom rules (verified end-to-end against Mooncake's rule test machinery), so the extension only adds zero-derivative rules for the functions that query shape/structure information and are not differentiable with respect to array contents. Co-Authored-By: Claude Fable 5 --- Project.toml | 3 ++ ext/ArraysOfArraysMooncakeExt.jl | 26 +++++++++++++++++ test/Project.toml | 2 ++ test/mooncake.jl | 49 ++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 81 insertions(+) create mode 100644 ext/ArraysOfArraysMooncakeExt.jl create mode 100644 test/mooncake.jl diff --git a/Project.toml b/Project.toml index e11f190..f082403 100644 --- a/Project.toml +++ b/Project.toml @@ -9,18 +9,21 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [extensions] ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" +ArraysOfArraysMooncakeExt = "Mooncake" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" [compat] Adapt = "1, 2, 3, 4" ChainRulesCore = "1" FixedSizeArrays = "1" +Mooncake = "0.5" StaticArraysCore = "1" Statistics = "1" julia = "1.10" diff --git a/ext/ArraysOfArraysMooncakeExt.jl b/ext/ArraysOfArraysMooncakeExt.jl new file mode 100644 index 0000000..a97ddb2 --- /dev/null +++ b/ext/ArraysOfArraysMooncakeExt.jl @@ -0,0 +1,26 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysMooncakeExt + +using Mooncake: DefaultCtx, @zero_derivative + +using ArraysOfArrays: ArraysOfArrays, AbstractSplitMode, + getsplitmode, unstackmode, is_memordered_splitmode, innersize, getslicemap, + consgrouped_ptrs + +# Mooncake can differentiate through the ArraysOfArrays types and operations +# without custom rules, so only the functions that query shape/structure +# information (and so are not differentiable with respect to array contents) +# get explicit zero-derivative rules: + +@zero_derivative DefaultCtx Tuple{typeof(getsplitmode), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(unstackmode), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(is_memordered_splitmode), AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(innersize), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(innersize), AbstractArray, Integer} +@zero_derivative DefaultCtx Tuple{typeof(getslicemap), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(ArraysOfArrays.getinnerdims), Tuple{Vararg{Integer}}, AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(ArraysOfArrays.getouterdims), Tuple{Vararg{Integer}}, AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(consgrouped_ptrs), AbstractVector} + +end # module ArraysOfArraysMooncakeExt diff --git a/test/Project.toml b/test/Project.toml index 221c071..6d81bf0 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -13,3 +14,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Documenter = "1" +Mooncake = "0.5" diff --git a/test/mooncake.jl b/test/mooncake.jl new file mode 100644 index 0000000..0ba665e --- /dev/null +++ b/test/mooncake.jl @@ -0,0 +1,49 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +import Mooncake +import Random + +@testset "Mooncake extension" begin + rng = Random.Xoshiro(0x2b6ee80184a75b47) + + x_flat = rand(rng, 3, 4) + x_vec = rand(rng, 10) + lengths = [2, 3, 5] + + @testset "zero-derivative primitives" begin + A = sliced(x_flat, Val(1)) + p = partitioned(x_vec, lengths) + + Mooncake.TestUtils.test_rule(rng, getsplitmode, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, getsplitmode, p; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, unstackmode, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, is_memordered_splitmode, getsplitmode(A); is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, innersize, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, getslicemap, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, consgrouped_ptrs, [1, 1, 2, 3, 3]; is_primitive = true, unsafe_perturb = true) + end + + # Mooncake differentiates the ArraysOfArrays types and operations without + # custom rules, verify gradient correctness end-to-end: + @testset "derived rules" begin + for (label, f, x) in [ + ("sliced and fused", x -> sum(abs2, fused(sliced(x, Val(1)))), x_flat), + ("stacked", x -> sum(stacked(ArrayOfSimilarArrays{Float64,1,1}(x))), x_flat), + ("vecflattened of ArrayOfSimilarArrays", x -> sum(abs2, vecflattened(ArrayOfSimilarArrays{Float64,1,1}(x))), x_flat), + ("partitioned", x -> sum(sum, partitioned(x, lengths)), x_vec), + ("VectorOfArrays round trip", x -> begin + p = partitioned(x, lengths) + sum(abs2, fused(splitup(fused(p), getsplitmode(p)))) + end, x_vec), + ("reduce vcat", x -> sum(sum, reduce(vcat, [partitioned(x, [2, 3]), partitioned(x, [4])])), x_vec), + ("innermap", x -> sum(fused(innermap(abs2, sliced(x, Val(1))))), x_flat), + ] + @testset "$label" begin + Mooncake.TestUtils.test_rule(rng, f, x; is_primitive = false, unsafe_perturb = true) + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index fc21dcc..6e66126 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,5 +11,6 @@ Test.@testset "Package ArraysOfArrays" begin include("broadcasting.jl") include("static_arrays.jl") include("fixed_size_arrays.jl") + include("mooncake.jl") include("test_docs.jl") end # testset From 009c6ecf59715275615c00d05971a4652a32587b Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 08:36:12 +0200 Subject: [PATCH 36/45] Make VectorOfArrays consistency checks GPU-compatible The consistency checks performed scalar indexing on elem_ptr, so VectorOfArrays with GPU-resident shape information could only be constructed with checks disabled. The checks now route through two helper functions that a new GPUArraysCore extension specializes: the two O(1) boundary reads run under @allowscalar and the O(n) partition check uses a vectorized formulation. Co-Authored-By: Claude Fable 5 --- Project.toml | 3 +++ ext/ArraysOfArraysGPUArraysCoreExt.jl | 31 +++++++++++++++++++++++++++ src/vector_of_arrays.jl | 24 +++++++++++++++------ 3 files changed, 51 insertions(+), 7 deletions(-) create mode 100644 ext/ArraysOfArraysGPUArraysCoreExt.jl diff --git a/Project.toml b/Project.toml index f082403..164806c 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" @@ -16,6 +17,7 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" +ArraysOfArraysGPUArraysCoreExt = "GPUArraysCore" ArraysOfArraysMooncakeExt = "Mooncake" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" @@ -23,6 +25,7 @@ ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" Adapt = "1, 2, 3, 4" ChainRulesCore = "1" FixedSizeArrays = "1" +GPUArraysCore = "0.2" Mooncake = "0.5" StaticArraysCore = "1" Statistics = "1" diff --git a/ext/ArraysOfArraysGPUArraysCoreExt.jl b/ext/ArraysOfArraysGPUArraysCoreExt.jl new file mode 100644 index 0000000..49fbf7d --- /dev/null +++ b/ext/ArraysOfArraysGPUArraysCoreExt.jl @@ -0,0 +1,31 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysGPUArraysCoreExt + +using GPUArraysCore: AbstractGPUArray, @allowscalar + +import ArraysOfArrays + +# Two O(1) scalar reads during construction are acceptable: +ArraysOfArrays._scalar_first_last(x::AbstractGPUArray) = @allowscalar (first(x), last(x)) + +# Vectorized formulation of the O(n) partition checks, to avoid scalar +# indexing of GPU arrays: +function ArraysOfArrays._partition_sizes_valid(elem_ptr::AbstractGPUArray{<:Integer}, kernel_size::AbstractVector) + ep_lo = view(elem_ptr, firstindex(elem_ptr):(lastindex(elem_ptr) - 1)) + ep_hi = view(elem_ptr, (firstindex(elem_ptr) + 1):lastindex(elem_ptr)) + len = ep_hi .- ep_lo + + klen = prod.(kernel_size) + # kernel_size may still live on the host: + klen_dev = klen isa AbstractGPUArray ? klen : copyto!(similar(elem_ptr, eltype(klen), size(klen)), klen) + + valid = (len .>= 0) .& ( + (klen_dev .== 1) .| ( + (klen_dev .!= 0) .& (mod.(len, max.(klen_dev, one(eltype(klen_dev)))) .== 0) + ) + ) + return all(valid) +end + +end # module ArraysOfArraysGPUArraysCoreExt diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 1b053f8..d9771fc 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -139,24 +139,34 @@ element_ptr(A::VectorOfArrays) = deepcopy(internal_element_ptr(A)) function full_consistency_checks(A::VectorOfArrays) simple_consistency_checks(A) - all(eachindex(A.kernel_size)) do i - len = A.elem_ptr[i+1] - A.elem_ptr[i] - klen = prod(A.kernel_size[i]) - len >= 0 && (klen == 1 || mod(len, klen) == 0) - end || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent")) + _partition_sizes_valid(A.elem_ptr, A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent")) nothing end +# Package extensions specialize this for GPU arrays, to avoid scalar indexing: +function _partition_sizes_valid(elem_ptr::AbstractVector{<:Integer}, kernel_size::AbstractVector) + all(eachindex(kernel_size)) do i + len = elem_ptr[i+1] - elem_ptr[i] + klen = prod(kernel_size[i]) + len >= 0 && (klen == 1 || mod(len, klen) == 0) + end +end + function simple_consistency_checks(A::VectorOfArrays{T,N,M}) where {T,N,M} M == N - 1 || throw(ArgumentError("VectorOfArrays{T,N,M} inconsistent: M must equal N - 1")) firstindex(A.elem_ptr) == firstindex(A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible indexing")) size(A.elem_ptr, 1) == size(A.kernel_size, 1) + 1 || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible size")) - first(A.elem_ptr) >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices")) - last(A.elem_ptr) - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices")) + ep_first, ep_last = _scalar_first_last(A.elem_ptr) + ep_first >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices")) + ep_last - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices")) nothing end +# Package extensions specialize this for GPU arrays, to allow the two scalar +# reads explicitly: +_scalar_first_last(x::AbstractVector) = (first(x), last(x)) + function no_consistency_checks(A::VectorOfArrays) nothing From 31e8810c37a5f374184f95a28250bac653553a08 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 08:36:12 +0200 Subject: [PATCH 37/45] Add KernelAbstractions get_backend for nested arrays Kernels over an ArrayOfSimilarArrays or VectorOfArrays run on the backend of the underlying data. The shape information of a VectorOfArrays may deliberately live on the host while its data lives on a device, so no cross-field consistency is required. Co-Authored-By: Claude Fable 5 --- Project.toml | 3 +++ ext/ArraysOfArraysKernelAbstractionsExt.jl | 16 ++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 ext/ArraysOfArraysKernelAbstractionsExt.jl diff --git a/Project.toml b/Project.toml index 164806c..4b9a855 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" @@ -18,6 +19,7 @@ ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" ArraysOfArraysGPUArraysCoreExt = "GPUArraysCore" +ArraysOfArraysKernelAbstractionsExt = "KernelAbstractions" ArraysOfArraysMooncakeExt = "Mooncake" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" @@ -26,6 +28,7 @@ Adapt = "1, 2, 3, 4" ChainRulesCore = "1" FixedSizeArrays = "1" GPUArraysCore = "0.2" +KernelAbstractions = "0.9" Mooncake = "0.5" StaticArraysCore = "1" Statistics = "1" diff --git a/ext/ArraysOfArraysKernelAbstractionsExt.jl b/ext/ArraysOfArraysKernelAbstractionsExt.jl new file mode 100644 index 0000000..7bedb5f --- /dev/null +++ b/ext/ArraysOfArraysKernelAbstractionsExt.jl @@ -0,0 +1,16 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysKernelAbstractionsExt + +import KernelAbstractions as KA + +using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays + +# Kernels over nested arrays run on the backend of the underlying data. The +# shape information of a VectorOfArrays may deliberately live on the host +# while its data lives on a device, so no cross-field consistency is +# required here: +KA.get_backend(A::ArrayOfSimilarArrays) = KA.get_backend(A.data) +KA.get_backend(A::VectorOfArrays) = KA.get_backend(A.data) + +end # module ArraysOfArraysKernelAbstractionsExt From 766cecb13962c1e26feed4d9e33dfdfd23b2cd39 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 08:36:41 +0200 Subject: [PATCH 38/45] Add GPU array tests and GPU usage docs Test both nested array types over JLArrays (the reference AbstractGPUArray implementation) with scalar indexing disallowed, and document the host-side and device-side VectorOfArrays layouts. Co-Authored-By: Claude Fable 5 --- docs/src/index.md | 4 +++ test/Project.toml | 3 ++ test/gpu_arrays.jl | 80 ++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 88 insertions(+) create mode 100644 test/gpu_arrays.jl diff --git a/docs/src/index.md b/docs/src/index.md index 3a12764..dd0335a 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -22,6 +22,10 @@ This package also defines and exports the following new functions applicable to All four are zero-copy where possible and so may return arrays that share memory with `A`. In contrast, `Base.stack(A)` and `reduce(vcat, A)` always return independent arrays. +## GPU support + +Both array types work with GPU-resident data. An `ArrayOfSimilarArrays` backed by a GPU array is fully functional. For a `VectorOfArrays`, the shape information (`elem_ptr` and `kernel_size`) can either stay on the host — element access then returns device views — or live on the device as well (e.g. via `Adapt.adapt`), which is the layout to use inside GPU kernels. `KernelAbstractions.get_backend` returns the backend of the underlying data. + ## [ArrayOfSimilarArrays](@id section_ArrayOfSimilarArrays) diff --git a/test/Project.toml b/test/Project.toml index 6d81bf0..4f12763 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,9 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/test/gpu_arrays.jl b/test/gpu_arrays.jl new file mode 100644 index 0000000..82824f2 --- /dev/null +++ b/test/gpu_arrays.jl @@ -0,0 +1,80 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using Adapt +using GPUArraysCore: AbstractGPUArray +using JLArrays +import KernelAbstractions as KA + +include("testdefs.jl") + +# JLArray is the reference AbstractGPUArray implementation, tests must not +# perform scalar indexing on it: +JLArrays.allowscalar(false) + +@testset "GPU arrays (via JLArrays)" begin + @testset "ArrayOfSimilarArrays on device data" begin + x = jl(rand(Float32, 3, 4, 5)) + A = @inferred ArrayOfSimilarArrays{Float32,1,2}(x) + + @test A[2, 3] isa AbstractGPUArray + @test @inferred(flatview(A)) === x + @test @inferred(fused(A)) === x + @test splitup(fused(A), getsplitmode(A)) == A + @test @inferred(innersize(A)) == (3,) + @test collect(fused(innermap(abs, A))) == abs.(collect(x)) + @test stack(A) isa AbstractGPUArray + A[1, 1] = jl(rand(Float32, 3)) + + V = VectorOfSimilarVectors(jl(rand(Float32, 3, 5))) + @test sum(V) isa AbstractGPUArray + @test ArraysOfArrays.Statistics.mean(V) isa AbstractGPUArray + @test vcat(V, V) isa VectorOfSimilarVectors + @test fused(vcat(V, V)) isa AbstractGPUArray + + @test KA.get_backend(A) == KA.get_backend(x) + end + + @testset "VectorOfArrays, device data and host shape info" begin + xv = jl(rand(Float32, 10)) + V = VectorOfArrays(xv, [1, 3, 6, 11], [(), (), ()]) + @test V[2] isa AbstractGPUArray + @test @inferred(vecflattened(V)) isa AbstractGPUArray + @test reduce(vcat, [V, V]) isa VectorOfArrays + @test partitioned(xv, [2, 3, 5]) isa VectorOfArrays + @test KA.get_backend(V) == KA.get_backend(xv) + end + + @testset "VectorOfArrays, everything on device" begin + xv = jl(rand(Float32, 10)) + + # Consistency checks must work without scalar indexing: + V = VectorOfArrays(xv, jl([1, 3, 6, 11]), jl([(), (), ()])) + @test V isa VectorOfArrays + @test length(V) == 3 + @test fused(V) === xv + + V2 = VectorOfArrays(jl(rand(Float32, 12)), jl([1, 5, 13]), jl([(2,), (2,)])) + @test length(V2) == 2 + + # Inconsistent shape info must still be detected: + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 3, 12]), jl([(), ()])) + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 6, 3]), jl([(), ()])) + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 4, 11]), jl([(2,), (2,)])) + + # Split mode round trip on device shape info: + sm = getsplitmode(V) + @test splitup(fused(V), sm) isa VectorOfArrays + @test typeof(splitup(fused(V), sm)) == typeof(V) + + # adapt moves data and shape info: + V_host = VectorOfArrays([Float32[1, 2], Float32[3, 4, 5]]) + V_dev = adapt(JLArray, V_host) + @test V_dev.data isa AbstractGPUArray + @test V_dev.elem_ptr isa AbstractGPUArray + @test adapt(Array, V_dev) == V_host + @test adapt(Array, getsplitmode(V_dev)).elem_ptr == getsplitmode(V_host).elem_ptr + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 6e66126..2ff6a88 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ Test.@testset "Package ArraysOfArrays" begin include("broadcasting.jl") include("static_arrays.jl") include("fixed_size_arrays.jl") + include("gpu_arrays.jl") include("mooncake.jl") include("test_docs.jl") end # testset From 1506ed0bd1e7efaae534298bf262bc4a14a18ec5 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 09:03:52 +0200 Subject: [PATCH 39/45] Add missing Random test dependency The Mooncake tests use Random, which was only available transitively. Co-Authored-By: Claude Fable 5 --- test/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/Project.toml b/test/Project.toml index 4f12763..14704e4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -10,6 +10,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" From 087c463645e2b9dbed05152996690000249864a4 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 15:32:42 +0200 Subject: [PATCH 40/45] Add mapat, innersizes and innerlengths mapat(f, Val(depth), As...) generalizes map/innermap/deepmap to arbitrary nesting depth and multiple (structure-aligned) inputs, similar to depth/axis-targeted operations in Python's AwkwardArrays, but with array-of-arrays nesting semantics. For split arrays it descends the split-mode stack and operates on the underlying flat data, in a single (GPU-compatible) operation per nesting level. innersizes and innerlengths return the per-element sizes/lengths of nested arrays without requiring equal element sizes, with vectorized (GPU-compatible) implementations for VectorOfArrays. Co-Authored-By: Claude Fable 5 --- src/functions.jl | 69 +++++++++++++++++++++++++++++++++ src/vector_of_arrays.jl | 22 +++++++---- test/array_of_similar_arrays.jl | 9 +++++ test/functions.jl | 22 +++++++++++ test/gpu_arrays.jl | 6 +++ test/vector_of_arrays.jl | 9 +++++ 6 files changed, 130 insertions(+), 7 deletions(-) diff --git a/src/functions.jl b/src/functions.jl index 7b15a85..80e9dae 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -356,6 +356,35 @@ end @inline innersize(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = getinnerdims(size(fused(A)), getsplitmode(A)) +""" + innersizes(A::AbstractArray{<:AbstractArray{T,M}}) + +Returns the sizes of the element arrays of `A`, as an array of `Dims{M}` +shaped like `A`. In contrast to [`innersize`](@ref), the element arrays do +not need to be of equal size. +""" +function innersizes end +export innersizes + +innersizes(A::AbstractArray{<:AbstractArray}) = map(size, A) + +innersizes(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = fill(innersize(A), size(A)) + + +""" + innerlengths(A::AbstractArray{<:AbstractArray}) + +Returns the lengths of the element arrays of `A`, as an array of `Int` +shaped like `A`. +""" +function innerlengths end +export innerlengths + +innerlengths(A::AbstractArray{<:AbstractArray}) = map(length, A) + +innerlengths(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = fill(prod(innersize(A)), size(A)) + + """ innermap(f, A::AbstractArray) innermap(f, A::AbstractArray{<:AbstractArray}) @@ -453,6 +482,46 @@ function _generic_deepmap_impl(f, A::AbstractArray) end +""" + mapat(f, ::Val{depth}, As::AbstractArray...) + +Nested `map` at nesting depth `depth`: apply `f` elementwise to the objects +at depth `depth` of the nested arrays `As`. Depth 1 refers to the elements +of the arrays themselves, so `mapat(f, Val(1), As...)` is equivalent to +`map(f, As...)` and `mapat(f, Val(2), A)` is equivalent to +[`innermap`](@ref)`(f, A)`. If `depth` exceeds the nesting depth of the +arrays, `f` is applied to the innermost elements, like [`deepmap`](@ref). + +All of `As` must have the same nesting structure down to `depth`, their +split modes (see [`getsplitmode`](@ref)) must be equal on each nesting +level. + +For split arrays (like [`ArrayOfSimilarArrays`](@ref) and +[`VectorOfArrays`](@ref)) `mapat` operates on the underlying flat data +directly, in a single (GPU-compatible) operation per nesting level. + +See also [`bcastat`](@ref) for broadcast semantics with arguments of +different nesting depth. +""" +function mapat end +export mapat + +@inline mapat(f, ::Val{1}, As::Vararg{AbstractArray,NA}) where {NA} = map(f, As...) + +function mapat(f, ::Val{depth}, As::Vararg{AbstractArray,NA}) where {depth,NA} + depth isa Integer && depth >= 1 || throw(ArgumentError("mapat depth must be a positive integer")) + smodes = map(getsplitmode, As) + smode = first(smodes) + if smode isa UnknownSplitMode + all(m -> m isa UnknownSplitMode, smodes) || throw(DimensionMismatch("mapat requires arrays with equal nesting structure, but split modes differ")) + return map((xs...) -> mapat(f, Val(depth - 1), xs...), As...) + else + all(isequal(smode), smodes) || throw(DimensionMismatch("mapat requires arrays with equal nesting structure, but split modes differ")) + return splitup(mapat(f, Val(depth - 1), map(fused, As)...), smode) + end +end + + """ partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) partitioned(A::AbstractVector, shapes::AbstractVector{<:Dims}) diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index d9771fc..9de9443 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -173,20 +173,28 @@ function no_consistency_checks(A::VectorOfArrays) end -Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer) - r = _voa_elem_range(A.elem_ptr, i) - len = length(r) - - ksize = A.kernel_size[i] +@inline function _elem_size(ksize::Dims, len::Integer) klen = prod(ksize) len_p, klen_p = promote(len, klen) - sz_lastdim = len == 0 ? len_p : div(len_p, klen_p) - sz = (ksize..., Int(sz_lastdim)) + sz_lastdim = len == 0 ? zero(len_p) : div(len_p, klen_p) + return (ksize..., Int(sz_lastdim)) +end +Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer) + r = _voa_elem_range(A.elem_ptr, i) + sz = _elem_size(A.kernel_size[i], length(r)) (r, sz) end +function innerlengths(A::VectorOfArrays) + ep = A.elem_ptr + return view(ep, firstindex(ep)+1:lastindex(ep)) .- view(ep, firstindex(ep):lastindex(ep)-1) +end + +innersizes(A::VectorOfArrays) = _elem_size.(A.kernel_size, innerlengths(A)) + + # Equality must be equivalent to elementwise comparison, but can be checked # much more efficiently via the flat data. The underlying data of A and B may # be offset relative to their element pointers, so only element sizes and the diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index fd229e0..967a068 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -174,6 +174,15 @@ end V = VectorOfSimilarVectors(V_flat) test_api(V, Array(V), V_flat) + # mapat operates on the flat data and preserves structure: + @test @inferred(mapat(abs2, Val(2), A)) == innermap(abs2, A) + @test typeof(mapat(abs2, Val(2), A)) == typeof(A) + @test fused(@inferred(mapat(+, Val(2), A, A))) == 2 .* A_flat + @test_throws DimensionMismatch mapat(+, Val(2), A, ArrayOfSimilarArrays{Float64,3,2}(A_flat)) + + @test @inferred(innersizes(A)) == fill(innersize(A), size(A)) + @test @inferred(innerlengths(A)) == fill(prod(innersize(A)), size(A)) + # vecflattened rrule: A_rr = ArrayOfSimilarArrays{Float64,1,1}(rand(3, 4)) y, pb = rrule(vecflattened, A_rr) diff --git a/test/functions.jl b/test/functions.jl index 7a5af8f..eb7468e 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -75,6 +75,28 @@ include("testdefs.jl") @test @inferred(deepmap(length, A_3)) == [[fill(1, 3), fill(1, 3)], [fill(1, 2), fill(1, 2)]] end + @testset "mapat" begin + f = x -> x^2 + @test @inferred(mapat(f, Val(1), A_1)) == map(f, A_1) + @test @inferred(mapat(f, Val(2), A_2)) == innermap(f, A_2) + @test @inferred(mapat(f, Val(3), A_3)) == deepmap(f, A_3) + @test mapat(length, Val(1), A_2) == length.(A_2) + + # Depth exceeding the nesting depth applies at the innermost level: + @test mapat(f, Val(5), A_1) == map(f, A_1) + + # Multiple inputs, zipped like map: + @test @inferred(mapat(+, Val(1), A_1, A_1)) == 2 .* A_1 + @test mapat(+, Val(2), A_2, A_2) == innermap(x -> 2 * x, A_2) + end + + @testset "innersizes and innerlengths" begin + @test @inferred(innersizes(A_2)) == size.(A_2) + @test @inferred(innerlengths(A_2)) == length.(A_2) + @test innersizes(A_3) == size.(A_3) + @test innerlengths(A_2e) == Int[] + end + @testset "Nested array API" begin test_api(A_0, A_0, A_0_flat) test_api(A_1, A_1, A_1_flat) diff --git a/test/gpu_arrays.jl b/test/gpu_arrays.jl index 82824f2..8f511d1 100644 --- a/test/gpu_arrays.jl +++ b/test/gpu_arrays.jl @@ -64,6 +64,12 @@ JLArrays.allowscalar(false) @test_throws ArgumentError VectorOfArrays(xv, jl([1, 6, 3]), jl([(), ()])) @test_throws ArgumentError VectorOfArrays(xv, jl([1, 4, 11]), jl([(2,), (2,)])) + # mapat and innerlengths only touch device data and shape info: + @test collect(fused(mapat(abs2, Val(2), V))) == abs2.(collect(xv)) + @test innerlengths(V) isa AbstractGPUArray + @test collect(innerlengths(V)) == [2, 3, 5] + @test collect(innersizes(V2)) == [(2, 2), (2, 4)] + # Split mode round trip on device shape info: sm = getsplitmode(V) @test splitup(fused(V), sm) isa VectorOfArrays diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index e247beb..de89bf3 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -209,6 +209,15 @@ include("testdefs.jl") test_api(B, Array(B), B.data) end + # mapat operates on the flat data and preserves structure: + @test @inferred(mapat(abs2, Val(2), B3)) == innermap(abs2, B3) + @test typeof(mapat(abs2, Val(2), B3)) == typeof(B3) + @test mapat(+, Val(2), B1, B1) == [x .+ x for x in B1] + @test_throws DimensionMismatch mapat(+, Val(2), B1, B3) + + @test @inferred(innerlengths(B1)) == length.(collect(B1)) + @test @inferred(innersizes(B3)) == size.(collect(B3)) + # getsplitmode must not be affected by later resizing: B_grow = VectorOfArrays([[1, 2], [3, 4, 5]]) sm_grow = getsplitmode(B_grow) From 879470716950cba21c1ea51e0f74e9b44140e352 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 15:43:12 +0200 Subject: [PATCH 41/45] Add bcastat bcastat(f, Val(depth), args...) broadcasts f over the contents of nested arrays at a given nesting depth, with AwkwardArrays-like alignment but array-of-arrays nesting semantics: aligned nested arguments, one value per element of shallower levels (broadcast over everything below), and scalars. For split arrays each nesting level compiles to a single flat (GPU-compatible) broadcast, using a segmented gather to expand outer-aligned values over ragged element contents. Co-Authored-By: Claude Fable 5 --- src/functions.jl | 88 +++++++++++++++++++++++++++++++++ src/vector_of_arrays.jl | 18 +++++++ test/array_of_similar_arrays.jl | 7 +++ test/gpu_arrays.jl | 7 +++ test/vector_of_arrays.jl | 39 +++++++++++++++ 5 files changed, 159 insertions(+) diff --git a/src/functions.jl b/src/functions.jl index 80e9dae..41f99c1 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -522,6 +522,94 @@ function mapat(f, ::Val{depth}, As::Vararg{AbstractArray,NA}) where {depth,NA} end +""" + bcastat(f, ::Val{depth}, args...) + +Broadcast `f` over the contents of nested arrays at nesting depth `depth`, +with AwkwardArrays-like alignment, but with array-of-arrays nesting +semantics: + +* Nested (split) arrays with equal split modes are aligned at depth `depth`. +* Arrays that match the outer structure of a shallower nesting level + contribute one value per element of that level, broadcast over everything + below it. +* Scalars and `Ref`s broadcast over everything. + +`bcastat(f, Val(1), args...)` is equivalent to `broadcast(f, args...)`, and +like [`mapat`](@ref), a `depth` that exceeds the nesting depth of the +arguments applies `f` at the innermost level. + +Nested array arguments must be split arrays (like +[`ArrayOfSimilarArrays`](@ref) and [`VectorOfArrays`](@ref)): for these, +`bcastat` compiles to a single flat (GPU-compatible) broadcast per nesting +level. +""" +function bcastat end +export bcastat + +@inline bcastat(f, ::Val{1}, args...) = broadcast(f, args...) + +function bcastat(f, ::Val{depth}, args...) where {depth} + depth isa Integer && depth >= 1 || throw(ArgumentError("bcastat depth must be a positive integer")) + r = _bcast_ref(args...) + if r === nothing + # No nested arguments left, apply at the innermost level: + foreach(_require_known_mode, args) + return broadcast(f, args...) + else + smode, ref_flat = r + descended = map(arg -> _bcast_descend(arg, smode, ref_flat), args) + return splitup(bcastat(f, Val(depth - 1), descended...), smode) + end +end + +function _bcast_ref(args...) + for a in args + if a isa AbstractArray + m = getsplitmode(a) + (m isa AbstractSlicingMode || m isa AbstractPartMode) && return (m, fused(a)) + end + end + return nothing +end + +_require_known_mode(@nospecialize(x)) = nothing + +function _require_known_mode(x::AbstractArray) + getsplitmode(x) isa UnknownSplitMode && throw(ArgumentError("bcastat requires nested array arguments to be split arrays with a known split mode, like VectorOfArrays or ArrayOfSimilarArrays")) + nothing +end + +function _bcast_descend(x, smode::AbstractSplitMode, ref_flat::AbstractArray) + x isa AbstractArray || return x + ndims(x) == 0 && return x + m = getsplitmode(x) + if m isa UnknownSplitMode + throw(ArgumentError("bcastat requires nested array arguments to be split arrays with a known split mode, like VectorOfArrays or ArrayOfSimilarArrays")) + elseif m isa NonSplitMode + return _bcast_expand(x, smode, ref_flat) + else + isequal(m, smode) || throw(DimensionMismatch("bcastat requires nested array arguments with equal split modes")) + return fused(x) + end +end + +function _bcast_expand(x::AbstractArray, smode::AbstractSlicingMode{M,N}, ref_flat::AbstractArray) where {M,N} + if axes(x) == getouterdims(axes(ref_flat), smode) + is_memordered_splitmode(smode) || throw(ArgumentError("bcastat does not support outer-value broadcasting for slicings that are not in memory order")) + return reshape(x, ntuple(_ -> 1, Val(M))..., size(x)...) + elseif axes(x) == axes(ref_flat) + return x + else + throw(DimensionMismatch("bcastat argument shape matches neither the outer structure nor the flat data of the nested arguments")) + end +end + +function _bcast_expand(x::AbstractArray, smode::AbstractPartMode, ref_flat::AbstractArray) + throw(ArgumentError("bcastat does not support outer-value broadcasting for split mode $(nameof(typeof(smode)))")) +end + + """ partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) partitioned(A::AbstractVector, shapes::AbstractVector{<:Dims}) diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 9de9443..4ff8d6c 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -307,6 +307,24 @@ function splitup(A::AbstractVector, smode::SplitParts) VectorOfArrays(A, smode.elem_ptr, smode.kernel_size, simple_consistency_checks) end +function _bcast_expand(x::AbstractArray, smode::SplitParts, ref_flat::AbstractArray) + ep = smode.elem_ptr + n_parts = length(smode.kernel_size) + if x isa AbstractVector && length(x) == n_parts + # One value per part, broadcast over the contents of that part. The + # value of x used for data not covered by any part is arbitrary, + # such data does not contribute to the result: + idx = similar(ref_flat, Int) + idx .= firstindex(ref_flat):lastindex(ref_flat) + seg = clamp.(searchsortedlast.(Ref(ep), idx), firstindex(x), lastindex(x)) + return x[seg] + elseif x isa AbstractVector && axes(x) == axes(ref_flat) + return x + else + throw(DimensionMismatch("bcastat argument shape matches neither the outer structure nor the flat data of the nested arguments")) + end +end + function partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) elem_ptr = _elem_ptr_from_lengths(A, lengths) diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 967a068..e99ce52 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -183,6 +183,13 @@ end @test @inferred(innersizes(A)) == fill(innersize(A), size(A)) @test @inferred(innerlengths(A)) == fill(prod(innersize(A)), size(A)) + # bcastat with outer-aligned, scalar and flat-matching arguments: + w = rand(size(A)...) + r_bc = @inferred bcastat(+, Val(2), A, w) + @test r_bc isa ArrayOfSimilarArrays + @test collect(r_bc) == [A[i] .+ w[i] for i in eachindex(A, w)] + @test fused(bcastat(muladd, Val(2), A, 2.0, A_flat)) == muladd.(A_flat, 2.0, A_flat) + # vecflattened rrule: A_rr = ArrayOfSimilarArrays{Float64,1,1}(rand(3, 4)) y, pb = rrule(vecflattened, A_rr) diff --git a/test/gpu_arrays.jl b/test/gpu_arrays.jl index 8f511d1..21cd68f 100644 --- a/test/gpu_arrays.jl +++ b/test/gpu_arrays.jl @@ -70,6 +70,13 @@ JLArrays.allowscalar(false) @test collect(innerlengths(V)) == [2, 3, 5] @test collect(innersizes(V2)) == [(2, 2), (2, 4)] + # bcastat compiles to flat device broadcasts: + v_outer = jl(Float32[10, 20, 30]) + r_bc = bcastat(+, Val(2), V, v_outer) + @test r_bc isa VectorOfArrays + @test fused(r_bc) isa AbstractGPUArray + @test collect(fused(r_bc)) == collect(xv) .+ [10, 10, 20, 20, 20, 30, 30, 30, 30, 30] + # Split mode round trip on device shape info: sm = getsplitmode(V) @test splitup(fused(V), sm) isa VectorOfArrays diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index de89bf3..e58fce3 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -272,6 +272,45 @@ include("testdefs.jl") end + @testset "bcastat" begin + x = collect(Float32, 1:10) + p = partitioned(x, [2, 3, 5]) + v = Float32[10, 20, 30] + + # One value per element, broadcast over the element contents: + r = @inferred bcastat(+, Val(2), p, v) + @test r isa PartsView + @test r == [xi .+ vi for (xi, vi) in zip(p, v)] + + # Scalars broadcast over everything: + @test bcastat(+, Val(2), p, 1) == [xi .+ 1 for xi in p] + + # Aligned nested arguments and flat-matching arguments: + @test bcastat(+, Val(2), p, p) == [xi .+ xi for xi in p] + @test bcastat(+, Val(2), p, x) == [xi .+ xi for xi in p] + + # Mixed argument kinds in one call: + @test bcastat(muladd, Val(2), p, v, 2) == [muladd.(xi, vi, 2) for (xi, vi) in zip(p, v)] + + # Data not covered by any part gets no contribution: + p_part = partitioned(x, [2, 3]) + @test bcastat(+, Val(2), p_part, Float32[10, 20]) == [Float32[11, 12], Float32[23, 24, 25]] + + # Depth exceeding the nesting depth applies at the innermost level: + @test bcastat(+, Val(3), p, v) == bcastat(+, Val(2), p, v) + + # Two nesting levels over a single flat buffer: + VV = VectorOfArrays(partitioned(collect(1.0:10.0), [2, 3, 5]), [1, 3, 4], [(), ()]) + w = [100.0, 200.0] + r2 = bcastat(+, Val(3), VV, w) + @test r2 isa VectorOfArrays + @test fused(fused(r2)) == [101, 102, 103, 104, 105, 206, 207, 208, 209, 210] + + @test_throws DimensionMismatch bcastat(+, Val(2), p, Float32[1, 2]) + @test_throws DimensionMismatch bcastat(+, Val(2), p, partitioned(x, [4, 6])) + @test_throws ArgumentError bcastat(+, Val(2), [[1, 2], [3]], 1) + end + @testset "rrules" begin x = collect(Float64, 1:10) From a6fe0212e859102664e1d68335f0b192d4781b54 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 15:54:29 +0200 Subject: [PATCH 42/45] Add innermapreduce, innerreduce and innersum Per-element reductions over the contents of nested arrays. For ArrayOfSimilarArrays this reduces over the inner dimensions of the flat data, for GPU-resident VectorOfArrays a new GPUArraysCore-plus- KernelAbstractions extension runs a single-pass segmented reduction kernel (moving host shape information to the device as required). Co-Authored-By: Claude Fable 5 --- Project.toml | 1 + ext/ArraysOfArraysGPUKernelsExt.jl | 65 ++++++++++++++++++++++++++++++ src/array_of_similar_arrays.jl | 13 ++++++ src/functions.jl | 50 +++++++++++++++++++++++ src/util.jl | 4 ++ test/array_of_similar_arrays.jl | 5 +++ test/gpu_arrays.jl | 9 +++++ test/vector_of_arrays.jl | 15 +++++++ 8 files changed, 162 insertions(+) create mode 100644 ext/ArraysOfArraysGPUKernelsExt.jl diff --git a/Project.toml b/Project.toml index 4b9a855..0b358a1 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" ArraysOfArraysGPUArraysCoreExt = "GPUArraysCore" +ArraysOfArraysGPUKernelsExt = ["GPUArraysCore", "KernelAbstractions"] ArraysOfArraysKernelAbstractionsExt = "KernelAbstractions" ArraysOfArraysMooncakeExt = "Mooncake" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" diff --git a/ext/ArraysOfArraysGPUKernelsExt.jl b/ext/ArraysOfArraysGPUKernelsExt.jl new file mode 100644 index 0000000..05c8862 --- /dev/null +++ b/ext/ArraysOfArraysGPUKernelsExt.jl @@ -0,0 +1,65 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysGPUKernelsExt + +using GPUArraysCore: AbstractGPUArray + +import KernelAbstractions as KA +using KernelAbstractions: @kernel, @index, @Const + +using ArraysOfArrays: ArraysOfArrays, VectorOfArrays, innerlengths + + +@kernel function _segmented_mapreduce_kernel!(out, f, op, @Const(data), @Const(elem_ptr), init) + i = @index(Global, Linear) + j0 = elem_ptr[i] + j1 = elem_ptr[i + 1] - 1 + if init isa ArraysOfArrays._NoInit + # Empty element arrays have been excluded beforehand: + acc = f(data[j0]) + for j in (j0 + 1):j1 + acc = op(acc, f(data[j])) + end + else + acc = init + for j in j0:j1 + acc = op(acc, f(data[j])) + end + end + out[i] = acc +end + + +function _on_backend(backend, x::AbstractArray) + KA.get_backend(x) == backend && return x + y = KA.allocate(backend, eltype(x), size(x)) + copyto!(y, x) + return y +end + + +# Single-pass segmented reduction for GPU-resident data: +function ArraysOfArrays._innermapreduce_impl(f, op, init, A::VectorOfArrays{T,N,M,<:AbstractGPUArray}) where {T,N,M} + data = A.data + backend = KA.get_backend(data) + elem_ptr = _on_backend(backend, A.elem_ptr) + n = length(A) + + if init isa ArraysOfArrays._NoInit && any(iszero, innerlengths(A)) + throw(ArgumentError("Reducing over empty element arrays requires an init value")) + end + + T_f = Base.promote_op(f, T) + T_init = init isa ArraysOfArrays._NoInit ? T_f : typeof(init) + T_out = promote_type(T_f, Base.promote_op(op, T_init, T_f)) + out = KA.allocate(backend, T_out, n) + + if n > 0 + # Stream-ordered, like GPU array operations no synchronization here: + kernel! = _segmented_mapreduce_kernel!(backend) + kernel!(out, f, op, data, elem_ptr, init; ndrange = n) + end + return out +end + +end # module ArraysOfArraysGPUKernelsExt diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 349d655..87bb345 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -96,6 +96,19 @@ Base.stack(A::AbstractArrayOfSimilarArrays; dims::Union{Integer,Colon} = :) = _s _stack_impl(A::AbstractArrayOfSimilarArrays, ::Colon) = copy(fused(A)) _stack_impl(A::AbstractArrayOfSimilarArrays, dims::Integer) = stack(collect(A); dims) +# Per-element reductions reduce over the inner dimensions of the flat data: +function _innermapreduce_impl(f, op, init, A::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} + data = fused(A) + dims = ntuple(identity, Val(M)) + r = if init isa _NoInit + mapreduce(f, op, data; dims = dims) + else + mapreduce(f, op, data; dims = dims, init = init) + end + return reshape(r, size(A)) +end + + @inline Base.:(==)(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = (stacked(A) == stacked(B)) @inline Base.isequal(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = isequal(stacked(A), stacked(B)) @inline Base.isapprox(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}; kwargs...) where {M,N} = isapprox(stacked(A), stacked(B); kwargs...) diff --git a/src/functions.jl b/src/functions.jl index 41f99c1..e03aad6 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -522,6 +522,56 @@ function mapat(f, ::Val{depth}, As::Vararg{AbstractArray,NA}) where {depth,NA} end +""" + innermapreduce(f, op, A::AbstractArray{<:AbstractArray}; [init]) + +Per-element `mapreduce` over the contents of the element arrays of `A`: +returns an array shaped like `A` that contains +`mapreduce(f, op, A[i]; [init])` for each element `A[i]`. + +For split arrays (like [`ArrayOfSimilarArrays`](@ref) and +[`VectorOfArrays`](@ref)) this uses efficient (GPU-compatible) segmented +reductions over the underlying flat data where possible. +""" +function innermapreduce end +export innermapreduce + +innermapreduce(f, op, A::AbstractArray{<:AbstractArray}; init = _NoInit()) = _innermapreduce_impl(f, op, init, A) + +function _innermapreduce_impl(f, op, init, A::AbstractArray{<:AbstractArray}) + if init isa _NoInit + map(x -> mapreduce(f, op, x), A) + else + map(x -> mapreduce(f, op, x; init = init), A) + end +end + + +""" + innerreduce(op, A::AbstractArray{<:AbstractArray}; [init]) + +Per-element `reduce` over the contents of the element arrays of `A`, +equivalent to `innermapreduce(identity, op, A; [init])`. +""" +function innerreduce end +export innerreduce + +innerreduce(op, A::AbstractArray{<:AbstractArray}; init = _NoInit()) = _innermapreduce_impl(identity, op, init, A) + + +""" + innersum(A::AbstractArray{<:AbstractArray}) + +Per-element sum over the contents of the element arrays of `A`, empty +element arrays sum to zero. +""" +function innersum end +export innersum + +innersum(A::AbstractArray{<:AbstractArray}) = innermapreduce(identity, +, A) +innersum(A::AbstractArray{<:AbstractArray{T}}) where {T<:Number} = innermapreduce(identity, +, A, init = zero(T)) + + """ bcastat(f, ::Val{depth}, args...) diff --git a/src/util.jl b/src/util.jl index adde6d6..8245511 100644 --- a/src/util.jl +++ b/src/util.jl @@ -43,6 +43,10 @@ end Base.@pure _val_value(::Val{x}) where x = x +# Internal sentinel for "no init value given": +struct _NoInit end + + # Concatenate arrays of equal ndims and equal size except in their last # dimension along their last dimension, with a single allocation: _cat_lastdim(datas) = _cat_lastdim_impl(datas, Val(ndims(first(datas)))) diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index e99ce52..96ba40d 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -183,6 +183,11 @@ end @test @inferred(innersizes(A)) == fill(innersize(A), size(A)) @test @inferred(innerlengths(A)) == fill(prod(innersize(A)), size(A)) + # Per-element reductions over the inner dimensions of the flat data: + @test @inferred(innersum(A)) == [sum(x) for x in A] + @test @inferred(innermapreduce(abs2, +, A)) ≈ [sum(abs2, x) for x in A] + @test @inferred(innerreduce(max, A)) == [maximum(x) for x in A] + # bcastat with outer-aligned, scalar and flat-matching arguments: w = rand(size(A)...) r_bc = @inferred bcastat(+, Val(2), A, w) diff --git a/test/gpu_arrays.jl b/test/gpu_arrays.jl index 21cd68f..8603ad4 100644 --- a/test/gpu_arrays.jl +++ b/test/gpu_arrays.jl @@ -77,6 +77,15 @@ JLArrays.allowscalar(false) @test fused(r_bc) isa AbstractGPUArray @test collect(fused(r_bc)) == collect(xv) .+ [10, 10, 20, 20, 20, 30, 30, 30, 30, 30] + # Segmented reductions run as a single device kernel: + xv_h = collect(xv) + parts_h = [xv_h[1:2], xv_h[3:5], xv_h[6:10]] + @test innersum(V) isa AbstractGPUArray + @test collect(innersum(V)) ≈ sum.(parts_h) + @test collect(innermapreduce(abs2, +, V)) ≈ [sum(abs2, p) for p in parts_h] + @test collect(innerreduce(max, V; init = -Inf32)) ≈ maximum.(parts_h) + @test_throws ArgumentError innerreduce(max, VectorOfArrays(xv, jl([1, 3, 3, 11]), jl([(), (), ()]))) + # Split mode round trip on device shape info: sm = getsplitmode(V) @test splitup(fused(V), sm) isa VectorOfArrays diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index e58fce3..2f786ce 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -311,6 +311,21 @@ include("testdefs.jl") @test_throws ArgumentError bcastat(+, Val(2), [[1, 2], [3]], 1) end + @testset "inner reductions" begin + x = collect(Float32, 1:10) + p = partitioned(x, [2, 3, 5]) + + @test @inferred(innermapreduce(abs2, +, p)) == [sum(abs2, xi) for xi in p] + @test @inferred(innerreduce(max, p)) == [maximum(xi) for xi in p] + @test @inferred(innersum(p)) == [sum(xi) for xi in p] + + # Empty element arrays require an init value, except for innersum: + pe = partitioned(collect(Float32, 1:5), [2, 0, 3]) + @test innersum(pe) == Float32[3, 0, 12] + @test_throws ArgumentError innerreduce(max, pe) + @test innerreduce(max, pe; init = -Inf32) == Float32[2, -Inf32, 5] + end + @testset "rrules" begin x = collect(Float64, 1:10) From b381719383572fe9664a8302becf5f9b74efbaf5 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 16:03:38 +0200 Subject: [PATCH 43/45] Structure-preserving outer broadcasts for vectors of arrays VectorOfArrays and AbstractArrayOfSimilarArrays get a broadcast style: outer-level broadcasts over a single outer dimension whose results are arrays now return a VectorOfArrays (with contiguous storage) instead of a Vector of individual arrays. The semantics of broadcasting are unchanged, f still receives whole element arrays, and all other cases fall back to the default broadcast machinery. Co-Authored-By: Claude Fable 5 --- src/broadcasting.jl | 41 ++++++++++++++++++++++++++++++++++++ test/vector_of_arrays.jl | 45 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+) diff --git a/src/broadcasting.jl b/src/broadcasting.jl index 9da047e..95b83a6 100644 --- a/src/broadcasting.jl +++ b/src/broadcasting.jl @@ -4,6 +4,47 @@ const _RefLike{T} = Union{Tuple{T}, Ref{T}} +""" + ArraysOfArrays.NestedArrayStyle{N}() + +Broadcast style of nested array types like [`VectorOfArrays`](@ref) and +[`ArrayOfSimilarArrays`](@ref). + +Broadcasts over a single outer dimension that produce array-valued results +(`f` receives whole element arrays and returns arrays, as in +`(x -> 2 .* x).(A)`) return a [`VectorOfArrays`](@ref) instead of a `Vector` +of arrays, where possible. All other broadcasts behave like the default +broadcast machinery. Use [`bcastat`](@ref) to broadcast over the *contents* +of the element arrays instead. +""" +struct NestedArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end + +NestedArrayStyle{M}(::Val{N}) where {M,N} = NestedArrayStyle{N}() + +Base.Broadcast.BroadcastStyle(::Type{<:AbstractArrayOfSimilarArrays{<:Any,<:Any,N}}) where {N} = NestedArrayStyle{N}() +Base.Broadcast.BroadcastStyle(::Type{<:VectorOfArrays}) = NestedArrayStyle{1}() + +function Base.copy(bc::Broadcast.Broadcasted{NestedArrayStyle{N}}) where {N} + ElType = Broadcast.combine_eltypes(bc.f, bc.args) + if N == 1 && ElType <: Array && isconcretetype(ElType) && axes(bc, 1) isa Base.OneTo + return _collect_nested(bc, ElType) + else + # Everything else behaves like the default broadcast machinery: + return copy(convert(Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{N}}, bc)) + end +end + +function _collect_nested(bc::Broadcast.Broadcasted, ::Type{Array{T,M}}) where {T,M} + dest = VectorOfArrays{T,M}() + sizehint!(dest.elem_ptr, length(axes(bc, 1)) + 1) + sizehint!(dest.kernel_size, length(axes(bc, 1))) + for i in eachindex(bc) + push!(dest, bc[i]) + end + return dest +end + + _idx_type(::VectorOfSimilarArrays) = Int _idx_type(A::VectorOfArrays) = eltype(A.elem_ptr) diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 2f786ce..58e9d9b 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -311,6 +311,51 @@ include("testdefs.jl") @test_throws ArgumentError bcastat(+, Val(2), [[1, 2], [3]], 1) end + @testset "outer broadcast" begin + A = VectorOfArrays([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0]]) + A_ref = collect(A) + + # Array-valued results at the outer level yield a VectorOfArrays: + r = (x -> 2 .* x).(A) + @test r isa VectorOfArrays{Float64,1} + @test r == [2 .* x for x in A_ref] + full_consistency_checks(r) + + # Results may be ragged in new ways: + r2 = (x -> vcat(x, sum(x))).(A) + @test r2 isa VectorOfArrays{Float64,1} + @test r2 == [vcat(x, sum(x)) for x in A_ref] + + # Multiple and mixed arguments: + @test broadcast((x, y) -> x .+ y, A, A) == [x .+ x for x in A_ref] + @test broadcast((x, y) -> x .+ y, A, A) isa VectorOfArrays + @test broadcast((x, s) -> x .* s, A, 2.0) isa VectorOfArrays + @test broadcast((x, s) -> x .+ s, A, [10.0, 20.0, 30.0]) == [x .+ s for (x, s) in zip(A_ref, [10.0, 20.0, 30.0])] + + # Scalar-valued results stay plain arrays: + @test sum.(A) isa Vector{Float64} + @test sum.(A) == sum.(A_ref) + + # Matrix elements: + B = VectorOfArrays([rand(2, 2), rand(3, 2)]) + rB = (x -> 2 .* x).(B) + @test rB isa VectorOfArrays{Float64,2} + @test rB == [2 .* x for x in collect(B)] + + # VectorOfSimilarArrays broadcasts to VectorOfArrays too: + C = VectorOfSimilarVectors(rand(3, 4)) + rC = (x -> x .+ 1).(C) + @test rC isa VectorOfArrays{Float64,1} + @test rC == [x .+ 1 for x in collect(C)] + + # Multi-dimensional outer structure falls back to the default + # behavior: + D = ArrayOfSimilarArrays{Float64,1,2}(rand(2, 3, 4)) + rD = (x -> 2 .* x).(D) + @test rD isa Matrix{Vector{Float64}} + @test rD == [2 .* D[i, j] for i in axes(D, 1), j in axes(D, 2)] + end + @testset "inner reductions" begin x = collect(Float32, 1:10) p = partitioned(x, [2, 3, 5]) From 4a19b6f8581101f259d4661e239adb6312b256b7 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 16:04:10 +0200 Subject: [PATCH 44/45] Document depth-targeted operations Co-Authored-By: Claude Fable 5 --- docs/src/index.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/src/index.md b/docs/src/index.md index dd0335a..70a66c9 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -13,6 +13,17 @@ This package also defines and exports the following new functions applicable to * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. +## Operations at a given nesting depth + +Similar to axis-targeted operations in Python's AwkwardArrays, but with array-of-arrays nesting semantics: + +* [`mapat(f, Val(d), As...)`](@ref) maps `f` over the objects at nesting depth `d` (`d = 1` ≡ `map`, `d = 2` ≡ `innermap`). +* [`bcastat(f, Val(d), args...)`](@ref) broadcasts `f` at depth `d`: nested arguments align, shallower arrays contribute one value per element of their level, scalars broadcast over everything. +* [`innermapreduce`](@ref), [`innerreduce`](@ref) and [`innersum`](@ref) reduce over the contents of each element array. +* [`innersizes`](@ref) and [`innerlengths`](@ref) return per-element sizes/lengths (elements need not be of equal size). + +For split arrays these operate on the underlying flat data — a single (GPU-compatible) operation per nesting level. Outer-level broadcasts like `(x -> 2 .* x).(A)` keep their usual Julia semantics (`f` receives whole element arrays), but return a `VectorOfArrays` when the results are arrays. + ## Which flattening function do I want? * [`fused(A)`](@ref): the original underlying array, `splitup(fused(A), getsplitmode(A)) == A`. From fdc7329036ab152ded7f0f233ed694d5b181f3b8 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 17:36:56 +0200 Subject: [PATCH 45/45] Accept Integer depth in mapat and bcastat Convenience methods that wrap the depth in a Val immediately, relying on constant propagation for type stability (verified via inference tests through a function barrier), like sliced already does for the number of inner dimensions. Co-Authored-By: Claude Fable 5 --- src/functions.jl | 13 ++++++++++++- test/functions.jl | 4 ++++ test/vector_of_arrays.jl | 4 ++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/functions.jl b/src/functions.jl index e03aad6..7fccf6b 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -484,6 +484,7 @@ end """ mapat(f, ::Val{depth}, As::AbstractArray...) + mapat(f, depth::Integer, As::AbstractArray...) Nested `map` at nesting depth `depth`: apply `f` elementwise to the objects at depth `depth` of the nested arrays `As`. Depth 1 refers to the elements @@ -492,6 +493,9 @@ of the arrays themselves, so `mapat(f, Val(1), As...)` is equivalent to [`innermap`](@ref)`(f, A)`. If `depth` exceeds the nesting depth of the arrays, `f` is applied to the innermost elements, like [`deepmap`](@ref). +The `Integer` form relies on constant propagation for type stability, use +the `Val` form when passing a non-constant depth. + All of `As` must have the same nesting structure down to `depth`, their split modes (see [`getsplitmode`](@ref)) must be equal on each nesting level. @@ -506,6 +510,8 @@ different nesting depth. function mapat end export mapat +@inline mapat(f, depth::Integer, As::Vararg{AbstractArray,NA}) where {NA} = mapat(f, Val(depth), As...) + @inline mapat(f, ::Val{1}, As::Vararg{AbstractArray,NA}) where {NA} = map(f, As...) function mapat(f, ::Val{depth}, As::Vararg{AbstractArray,NA}) where {depth,NA} @@ -574,6 +580,7 @@ innersum(A::AbstractArray{<:AbstractArray{T}}) where {T<:Number} = innermapreduc """ bcastat(f, ::Val{depth}, args...) + bcastat(f, depth::Integer, args...) Broadcast `f` over the contents of nested arrays at nesting depth `depth`, with AwkwardArrays-like alignment, but with array-of-arrays nesting @@ -587,7 +594,9 @@ semantics: `bcastat(f, Val(1), args...)` is equivalent to `broadcast(f, args...)`, and like [`mapat`](@ref), a `depth` that exceeds the nesting depth of the -arguments applies `f` at the innermost level. +arguments applies `f` at the innermost level. The `Integer` form relies on +constant propagation for type stability, use the `Val` form when passing a +non-constant depth. Nested array arguments must be split arrays (like [`ArrayOfSimilarArrays`](@ref) and [`VectorOfArrays`](@ref)): for these, @@ -597,6 +606,8 @@ level. function bcastat end export bcastat +@inline bcastat(f, depth::Integer, args...) = bcastat(f, Val(depth), args...) + @inline bcastat(f, ::Val{1}, args...) = broadcast(f, args...) function bcastat(f, ::Val{depth}, args...) where {depth} diff --git a/test/functions.jl b/test/functions.jl index eb7468e..85612c9 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -88,6 +88,10 @@ include("testdefs.jl") # Multiple inputs, zipped like map: @test @inferred(mapat(+, Val(1), A_1, A_1)) == 2 .* A_1 @test mapat(+, Val(2), A_2, A_2) == innermap(x -> 2 * x, A_2) + + # Integer depth relies on constant propagation for type stability: + mapat_intdepth(g, A) = mapat(g, 2, A) + @test @inferred(mapat_intdepth(f, A_2)) == innermap(f, A_2) end @testset "innersizes and innerlengths" begin diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 58e9d9b..6862289 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -299,6 +299,10 @@ include("testdefs.jl") # Depth exceeding the nesting depth applies at the innermost level: @test bcastat(+, Val(3), p, v) == bcastat(+, Val(2), p, v) + # Integer depth relies on constant propagation for type stability: + bcastat_intdepth(g, A, y) = bcastat(g, 2, A, y) + @test @inferred(bcastat_intdepth(+, p, v)) == bcastat(+, Val(2), p, v) + # Two nesting levels over a single flat buffer: VV = VectorOfArrays(partitioned(collect(1.0:10.0), [2, 3, 5]), [1, 3, 4], [(), ()]) w = [100.0, 200.0]