From 2094a782cc93b08a20a5c8adca39abe3b5e35ddf Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 9 May 2025 07:55:38 -0400 Subject: [PATCH 1/6] Type-generic zeromatrix for arraypartition One part of https://github.com/SciML/OrdinaryDiffEq.jl/issues/2703 --- src/array_partition.jl | 5 ++++- test/gpu/arraypartition_gpu.jl | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index b0325fe0..33b328f2 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -430,7 +430,10 @@ end ## Linear Algebra -ArrayInterface.zeromatrix(A::ArrayPartition) = ArrayInterface.zeromatrix(Vector(A)) +function ArrayInterface.zeromatrix(A::ArrayPartition) + x = reduce(vcat,vec.(A.x)) + x .* x' .* false +end function __get_subtypes_in_module( mod, supertype; include_supertype = true, all = false, except = []) diff --git a/test/gpu/arraypartition_gpu.jl b/test/gpu/arraypartition_gpu.jl index 08fdb69a..e0914bac 100644 --- a/test/gpu/arraypartition_gpu.jl +++ b/test/gpu/arraypartition_gpu.jl @@ -1,4 +1,4 @@ -using RecursiveArrayTools, CUDA, Test +using RecursiveArrayTools, ArrayInterface, CUDA, Test CUDA.allowscalar(false) # Test indexing with colon @@ -21,3 +21,6 @@ fill!(pA, false) a = ArrayPartition(([1.0f0] |> cu, [2.0f0] |> cu, [3.0f0] |> cu)) b = ArrayPartition(([0.0f0] |> cu, [0.0f0] |> cu, [0.0f0] |> cu)) @. a + b + +@test ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2)))) isa CuMatrix +@test size(ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2))))) == (4,4) From cd9a9750e66eff1e150a4a8bfe254f2b4a34329e Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 9 May 2025 08:25:33 -0400 Subject: [PATCH 2/6] fix mapreduce recursion --- src/array_partition.jl | 2 +- test/gpu/arraypartition_gpu.jl | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index 33b328f2..39e28ffe 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -166,7 +166,7 @@ Base.:(==)(A::ArrayPartition, B::ArrayPartition) = A.x == B.x Base.map(f, A::ArrayPartition) = ArrayPartition(map(x -> map(f, x), A.x)) function Base.mapreduce(f, op, A::ArrayPartition{T}; kwargs...) where {T} - mapreduce(f, op, (i for i in A); kwargs...) + mapreduce(x->mapreduce(f, op, x; kwargs...), op, (i for i in A.x); kwargs...) end Base.filter(f, A::ArrayPartition) = ArrayPartition(map(x -> filter(f, x), A.x)) Base.any(f, A::ArrayPartition) = any((any(f, x) for x in A.x)) diff --git a/test/gpu/arraypartition_gpu.jl b/test/gpu/arraypartition_gpu.jl index e0914bac..2c457e79 100644 --- a/test/gpu/arraypartition_gpu.jl +++ b/test/gpu/arraypartition_gpu.jl @@ -22,5 +22,7 @@ a = ArrayPartition(([1.0f0] |> cu, [2.0f0] |> cu, [3.0f0] |> cu)) b = ArrayPartition(([0.0f0] |> cu, [0.0f0] |> cu, [0.0f0] |> cu)) @. a + b -@test ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2)))) isa CuMatrix -@test size(ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2))))) == (4,4) +x = ArrayPartition((CUDA.zeros(2),CUDA.zeros(2))) +@test ArrayInterface.zeromatrix(x) isa CuMatrix +@test size(ArrayInterface.zeromatrix(x)) == (4,4) +@test maximum(abs, x) == 0f0 From daa4cb27582080b085315ca2f07d2d7f4020f365 Mon Sep 17 00:00:00 2001 From: HenrySnowden Date: Mon, 12 May 2025 14:42:31 +0100 Subject: [PATCH 3/6] Added the zeromatrix function to NamedArrayPartition A function similar to what is implemented commit 2094a78 but for NamedArrayPartitions rather than ArrayPartitions. Tested privately to work with Implicit solvers in OrdinaryDiffEq.jl --- src/named_array_partition.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/named_array_partition.jl b/src/named_array_partition.jl index de8fa91a..99c83512 100644 --- a/src/named_array_partition.jl +++ b/src/named_array_partition.jl @@ -145,6 +145,13 @@ end return dest end +#Overwrite ArrayInterface zeromatrix to work with NamedArrayPartitions & implicit solvers within OrdinaryDiffEq +function ArrayInterface.zeromatrix(A::NamedArrayPartition) + B = ArrayPartition(A) + x = reduce(vcat,vec.(B.x)) + x .* x' .* false +end + # `x = find_NamedArrayPartition(x)` returns the first `NamedArrayPartition` among broadcast arguments. find_NamedArrayPartition(bc::Base.Broadcast.Broadcasted) = find_NamedArrayPartition(bc.args) function find_NamedArrayPartition(args::Tuple) From 3450c85ab59c9eb51115bd641fb229c42b21e6e3 Mon Sep 17 00:00:00 2001 From: HenrySnowden Date: Mon, 12 May 2025 14:50:09 +0100 Subject: [PATCH 4/6] Added tests for new zero matrix function --- test/named_array_partition_tests.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/named_array_partition_tests.jl b/test/named_array_partition_tests.jl index d5647bad..e6747969 100644 --- a/test/named_array_partition_tests.jl +++ b/test/named_array_partition_tests.jl @@ -9,10 +9,13 @@ using RecursiveArrayTools, Test @test x.a ≈ ones(10) @test typeof(x .+ x[1:end]) <: Vector # test broadcast precedence @test all(x .== x[1:end]) + @test ArrayInterface.zeromatrix(x) isa Matrix + @test size(ArrayInterface.zeromatrix(x)) == (30,30) y = copy(x) @test zero(x, (10, 20)) == zero(x) # test that ignoring dims works @test typeof(zero(x)) <: NamedArrayPartition @test (y .*= 2).a[1] ≈ 2 # test in-place bcast + @test length(Array(x)) == 30 @test typeof(Array(x)) <: Array From e11c4cdb7239cdf9051ee4e19ad2982c3075d509 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas Date: Mon, 8 Dec 2025 16:19:36 -0500 Subject: [PATCH 5/6] Fix mapreduce type-stability on Julia 1.10 using @generated functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Use @generated functions (_mapreduce_impl and _mapreduce_impl_init) to ensure type-stable mapreduce for ArrayPartition on Julia 1.10 - The generated approach unrolls the tuple iteration at compile time, avoiding type inference issues with kwargs that affected Julia 1.10 - Preserves correct `init` parameter semantics (init applied once at outer level) - Add missing ArrayInterface import in named_array_partition_tests.jl This fixes the type inference failure where Julia 1.10 would infer `Any` instead of the correct concrete return type for mapreduce operations on nested ArrayPartitions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- src/array_partition.jl | 36 +++++++++++++++++++++++++++-- test/named_array_partition_tests.jl | 2 +- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index 39e28ffe..baac99bd 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -165,8 +165,40 @@ Base.:(==)(A::ArrayPartition, B::ArrayPartition) = A.x == B.x ## Iterable Collection Constructs Base.map(f, A::ArrayPartition) = ArrayPartition(map(x -> map(f, x), A.x)) -function Base.mapreduce(f, op, A::ArrayPartition{T}; kwargs...) where {T} - mapreduce(x->mapreduce(f, op, x; kwargs...), op, (i for i in A.x); kwargs...) +# Use @generated function for type stability on Julia 1.10 +# The generated approach avoids type inference issues with kwargs in older Julia versions +@generated function _mapreduce_impl(f, op, A::ArrayPartition{T, S}) where {T, S} + N = length(S.parameters) + if N == 1 + return :(mapreduce(f, op, A.x[1])) + else + expr = :(mapreduce(f, op, A.x[$N])) + for i in (N - 1):-1:1 + expr = :(op(mapreduce(f, op, A.x[$i]), $expr)) + end + return expr + end +end +@generated function _mapreduce_impl_init(f, op, A::ArrayPartition{T, S}, init) where {T, S} + N = length(S.parameters) + if N == 1 + return :(mapreduce(f, op, A.x[1])) + else + expr = :(mapreduce(f, op, A.x[$N])) + for i in (N - 1):-1:1 + expr = :(op(mapreduce(f, op, A.x[$i]), $expr)) + end + # Apply init only at the outermost reduction + return :(op(init, $expr)) + end +end +@inline function Base.mapreduce(f, op, A::ArrayPartition; + init = Base._InitialValue(), kwargs...) + if init isa Base._InitialValue + _mapreduce_impl(f, op, A) + else + _mapreduce_impl_init(f, op, A, init) + end end Base.filter(f, A::ArrayPartition) = ArrayPartition(map(x -> filter(f, x), A.x)) Base.any(f, A::ArrayPartition) = any((any(f, x) for x in A.x)) diff --git a/test/named_array_partition_tests.jl b/test/named_array_partition_tests.jl index e6747969..069cd456 100644 --- a/test/named_array_partition_tests.jl +++ b/test/named_array_partition_tests.jl @@ -1,4 +1,4 @@ -using RecursiveArrayTools, Test +using RecursiveArrayTools, ArrayInterface, Test @testset "NamedArrayPartition tests" begin x = NamedArrayPartition(a = ones(10), b = rand(20)) From b9c30c86f0a0d6288ddcf3f86559e031b4573f34 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Sun, 14 Dec 2025 08:57:54 -0500 Subject: [PATCH 6/6] Update arraypartition_gpu.jl --- test/gpu/arraypartition_gpu.jl | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/test/gpu/arraypartition_gpu.jl b/test/gpu/arraypartition_gpu.jl index 625d5463..85b59893 100644 --- a/test/gpu/arraypartition_gpu.jl +++ b/test/gpu/arraypartition_gpu.jl @@ -22,6 +22,26 @@ a = ArrayPartition(([1.0f0] |> cu, [2.0f0] |> cu, [3.0f0] |> cu)) b = ArrayPartition(([0.0f0] |> cu, [0.0f0] |> cu, [0.0f0] |> cu)) @. a + b +# Test adapt from ArrayPartition with CuArrays to ArrayPartition with CPU arrays + +a = CuArray(Float64.([1., 2., 3., 4.])) +b = CuArray(Float64.([1., 2., 3., 4.])) +part_a_gpu = ArrayPartition(a, b) +part_a = adapt(Array{Float32}, part_a_gpu) + +c = Float32.([1., 2., 3., 4.]) +d = Float32.([1., 2., 3., 4.]) +part_b = ArrayPartition(c, d) + +@test part_a == part_b # Test equality + +for i in 1:length(part_a.x) + sub_a = part_a.x[i] + sub_b = part_b.x[i] + @test sub_a == sub_b # Test for value equality in sub-arrays + @test typeof(sub_a) === typeof(sub_b) # Test type equality +end + x = ArrayPartition((CUDA.zeros(2),CUDA.zeros(2))) @test ArrayInterface.zeromatrix(x) isa CuMatrix @test size(ArrayInterface.zeromatrix(x)) == (4,4)