From 24669540cac84b0ab97cc921882cb3def2727fa8 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 27 Apr 2026 11:06:38 -0400 Subject: [PATCH 1/2] Fix NamedArrayPartition Vector{Int} / UnitRange indexing (#583) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The invalidation-eliminating change in 05faa730 narrowed `getindex(::NamedArrayPartition, ...)` from a catch-all `args...` form to just `Int`. Anything else now falls through `AbstractArray`'s generic indexing, which routes via `similar(x, T, dims)`. That exposed a latent bug in the `similar` overloads: when `dims != size(A)`, `similar(::ArrayPartition, T, dims)` intentionally degrades to a plain `Vector` (a partition layout no longer fits the requested shape). The `NamedArrayPartition` overloads then tried to wrap that `Vector` via the inner constructor `NamedArrayPartition{T, A<:ArrayPartition{T}, NT}(::A, ::NT)`, which raised a MethodError. The catch-all `args...` getindex used to bypass `similar` entirely, so the bug stayed hidden until 05faa730. Make `similar(::NAP, dims)` and `similar(::NAP, T, dims)` mirror `ArrayPartition`'s own degradation: wrap when the inner result is an `ArrayPartition`, otherwise return the fallback array as-is. This restores `x[1:2]` and `x[[1, 4]]` for partial slices, keeps `x[1:end]` returning a `NamedArrayPartition` (the full-range slice still hits the ArrayPartition path), and reintroduces no method that catches a wider signature than necessary — `using RecursiveArrayTools` still produces 0 invalidation trees, verified with `SnoopCompile.invalidation_trees`. Closes #583. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.7 (1M context) --- src/named_array_partition.jl | 20 +++++++++------- test/named_array_partition_tests.jl | 36 +++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/src/named_array_partition.jl b/src/named_array_partition.jl index b2546334..771e570a 100644 --- a/src/named_array_partition.jl +++ b/src/named_array_partition.jl @@ -34,11 +34,16 @@ function Base.similar(A::NamedArrayPartition) ) end -# return ArrayPartition when possible, otherwise next best thing of the correct size +# return NamedArrayPartition when the requested dims still match the partition layout; +# otherwise fall back to the plain backing array of the correct size. ArrayPartition's +# own `similar(A, dims)` already does this degradation (it returns a Vector when +# `dims != size(A)`), and we simply propagate that result instead of trying to +# wrap a non-ArrayPartition in a NamedArrayPartition (which would hit the inner +# constructor signature `NamedArrayPartition(::A<:ArrayPartition, ::NamedTuple)`). function Base.similar(A::NamedArrayPartition, dims::NTuple{N, Int}) where {N} - return NamedArrayPartition( - similar(getfield(A, :array_partition), dims), getfield(A, :names_to_indices) - ) + inner = similar(getfield(A, :array_partition), dims) + inner isa ArrayPartition || return inner + return NamedArrayPartition(inner, getfield(A, :names_to_indices)) end # similar array partition of common type @@ -48,11 +53,10 @@ end ) end -# return ArrayPartition when possible, otherwise next best thing of the correct size function Base.similar(A::NamedArrayPartition, ::Type{T}, dims::NTuple{N, Int}) where {T, N} - return NamedArrayPartition( - similar(getfield(A, :array_partition), T, dims), getfield(A, :names_to_indices) - ) + inner = similar(getfield(A, :array_partition), T, dims) + inner isa ArrayPartition || return inner + return NamedArrayPartition(inner, getfield(A, :names_to_indices)) end # similar array partition with different types diff --git a/test/named_array_partition_tests.jl b/test/named_array_partition_tests.jl index 9a012635..a6cbc867 100644 --- a/test/named_array_partition_tests.jl +++ b/test/named_array_partition_tests.jl @@ -37,3 +37,39 @@ using RecursiveArrayTools, ArrayInterface, Test @test typeof((x -> x[1]).(x)) <: NamedArrayPartition @test typeof(map(x -> x[1], x)) <: NamedArrayPartition end + +# Regression test for https://github.com/SciML/RecursiveArrayTools.jl/issues/583: +# indexing a NamedArrayPartition with a UnitRange / Vector{Int} smaller than +# the whole array used to throw a MethodError because `similar(::NAP, T, dims)` +# tried to wrap a plain Vector (returned by `similar(::ArrayPartition, T, dims)` +# when `dims != size(A)`) in NamedArrayPartition's inner constructor, which +# requires an ArrayPartition. +@testset "NamedArrayPartition issue #583 indexing" begin + x = NamedArrayPartition(a = ones(2), b = 2 * ones(3)) + + # UnitRange that doesn't span the whole array: returns a plain Vector + @test x[1:2] == [1.0, 1.0] + @test x[1:2] isa Vector{Float64} + @test x[2:4] == [1.0, 2.0, 2.0] + + # Vector{Int} indexing + @test x[[1, 2]] == [1.0, 1.0] + @test x[[1, 4]] == [1.0, 2.0] + @test x[[1, 4]] isa Vector{Float64} + + # Existing behavior: full-range slice still preserves NamedArrayPartition + @test typeof(x[1:end]) <: NamedArrayPartition + + # `similar` with a non-matching dims tuple gives the backing-array fallback + @test similar(x, Float64, (2,)) isa Vector{Float64} + @test similar(x, (2,)) isa Vector{Float64} + # `similar` with matching dims preserves the NamedArrayPartition wrapper + @test similar(x, Float64, size(x)) isa NamedArrayPartition + @test similar(x, size(x)) isa NamedArrayPartition + + # Scalar indexing untouched + @test x[1] == 1.0 + @test x[3] == 2.0 + x[1] = 99.0 + @test x[1] == 99.0 +end From 9cbf6d0affa7fab7a7f015c8a7eb910ea26ceac3 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 28 Apr 2026 10:36:56 -0400 Subject: [PATCH 2/2] Make NamedArrayPartition slicing fully type-stable via Vector The first commit on this branch fixed `x[1:2]` and `x[[1, 4]]` by making `similar(::NAP, T, dims)` degrade to a Vector when `dims != size(A)`, mirroring ArrayPartition. That worked, but it left the indexing path inferring as `Union{NamedArrayPartition, Vector{Float64}}` because `similar(::ArrayPartition, T, dims)` itself is a Union (the `dims == size(A)` branch is a runtime check). Add a `_unsafe_getindex(::IndexStyle, ::NAP, I::Vararg{Union{Real, AbstractArray}})` shortcut that mirrors the one at array_partition.jl:317. Allocate the destination directly off the underlying first array and fill it with `Base._unsafe_getindex!`. The shortcut bypasses `similar` entirely for the indexing path, so `x[1:2]`, `x[[1, 4]]`, `x[1:length(x)]` all infer to a clean `Vector{Float64}`. Trade-off: this regresses the post-05faa730 test `typeof(x .+ x[1:end]) <: NamedArrayPartition` back to `<: Vector` (the v3 behavior). That test was added in 05faa730 alongside the invalidation cleanup, but its only effect was to mask the unstable small-Union return; restoring v3 semantics here gives full type stability and matches what ArrayPartition already does. Use `similar(x)` / `copy(x)` if you want a NamedArrayPartition back. The `similar(::NAP, dims)` and `similar(::NAP, T, dims)` overloads keep the graceful-degrade-to-Vector behavior from the previous commit, so direct `similar(x, T, (2,))` calls (e.g. from downstream library code) still work. `SnoopCompile.invalidation_trees(@snoop_invalidations using RecursiveArrayTools)` still reports 0 trees, full `Pkg.test()` passes, and the new regression test asserts type stability via `@inferred`. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.7 (1M context) --- src/named_array_partition.jl | 22 +++++++++++++++ test/named_array_partition_tests.jl | 42 +++++++++++++++++------------ 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/src/named_array_partition.jl b/src/named_array_partition.jl index 771e570a..887f658e 100644 --- a/src/named_array_partition.jl +++ b/src/named_array_partition.jl @@ -100,6 +100,28 @@ Base.length(x::NamedArrayPartition) = length(ArrayPartition(x)) # Use concrete index types to avoid invalidating AbstractArray's generic setindex!. Base.@propagate_inbounds Base.getindex(x::NamedArrayPartition, i::Int) = ArrayPartition(x)[i] Base.@propagate_inbounds Base.setindex!(x::NamedArrayPartition, v, i::Int) = (ArrayPartition(x)[i] = v) + +# Indexing with non-scalar indices (UnitRange, Vector{Int}, etc.) goes through +# AbstractArray's generic path, which routes via `similar(A, T, dims)`. NAP's +# `similar(::NAP, T, dims)` cannot in general produce a NamedArrayPartition for +# arbitrary `dims` (the partition layout is fixed by `names_to_indices`), so it +# falls back to a plain Vector — making the inferred return type a small Union. +# +# Mirror ArrayPartition's `_unsafe_getindex` shortcut at `array_partition.jl:317`: +# allocate the destination directly off the first underlying array and fill it +# via `_unsafe_getindex!`. The result is always a Vector for non-scalar indexing, +# so `x[I]` is type-stable. This matches the v3 indexing semantics (`x[1:end]` +# returns a `Vector`, not a `NamedArrayPartition`); use `similar(x)` / +# `copy(x)` if you want a NamedArrayPartition back. +Base.@propagate_inbounds function Base._unsafe_getindex( + ::IndexStyle, A::NamedArrayPartition, + I::Vararg{Union{Real, AbstractArray}, N} + ) where {N} + shape = Base.index_shape(I...) + dest = similar(getfield(A, :array_partition).x[1], shape) + Base._unsafe_getindex!(dest, A, I...) + return dest +end function Base.map(f, x::NamedArrayPartition) return NamedArrayPartition(map(f, ArrayPartition(x)), getfield(x, :names_to_indices)) end diff --git a/test/named_array_partition_tests.jl b/test/named_array_partition_tests.jl index a6cbc867..1cc9f952 100644 --- a/test/named_array_partition_tests.jl +++ b/test/named_array_partition_tests.jl @@ -7,7 +7,7 @@ using RecursiveArrayTools, ArrayInterface, Test @test typeof(similar(x)) <: NamedArrayPartition @test typeof(similar(x, Int)) <: NamedArrayPartition @test x.a ≈ ones(10) - @test typeof(x .+ x[1:end]) <: NamedArrayPartition # x[1:end] preserves type + @test typeof(x .+ x[1:end]) <: Vector # x[1:end] is a plain Vector (type-stable slicing) @test all(x .== x[1:end]) @test ArrayInterface.zeromatrix(x) isa Matrix @test size(ArrayInterface.zeromatrix(x)) == (30, 30) @@ -39,37 +39,45 @@ using RecursiveArrayTools, ArrayInterface, Test end # Regression test for https://github.com/SciML/RecursiveArrayTools.jl/issues/583: -# indexing a NamedArrayPartition with a UnitRange / Vector{Int} smaller than -# the whole array used to throw a MethodError because `similar(::NAP, T, dims)` -# tried to wrap a plain Vector (returned by `similar(::ArrayPartition, T, dims)` -# when `dims != size(A)`) in NamedArrayPartition's inner constructor, which -# requires an ArrayPartition. +# indexing a NamedArrayPartition with a UnitRange / Vector{Int} smaller than the +# whole array used to throw a MethodError because the AbstractArray indexing +# path called `similar(::NAP, T, dims)`, which tried to wrap a plain Vector +# (returned by `similar(::ArrayPartition, T, dims)` for `dims != size(A)`) in +# NamedArrayPartition's inner constructor, which requires an ArrayPartition. +# +# The `_unsafe_getindex(::IndexStyle, ::NAP, I...)` shortcut bypasses `similar` +# entirely, allocating a plain Vector destination directly. Slicing therefore +# always returns a Vector and is type-stable. @testset "NamedArrayPartition issue #583 indexing" begin x = NamedArrayPartition(a = ones(2), b = 2 * ones(3)) - # UnitRange that doesn't span the whole array: returns a plain Vector + # UnitRange / Vector{Int} indexing all return Vector and are type-stable @test x[1:2] == [1.0, 1.0] - @test x[1:2] isa Vector{Float64} @test x[2:4] == [1.0, 2.0, 2.0] - - # Vector{Int} indexing + @test x[1:end] == [1.0, 1.0, 2.0, 2.0, 2.0] @test x[[1, 2]] == [1.0, 1.0] @test x[[1, 4]] == [1.0, 2.0] + + @test x[1:2] isa Vector{Float64} + @test x[1:end] isa Vector{Float64} @test x[[1, 4]] isa Vector{Float64} - # Existing behavior: full-range slice still preserves NamedArrayPartition - @test typeof(x[1:end]) <: NamedArrayPartition + # Inferred return types: Vector, not Union + @test (@inferred x[1:2]) isa Vector{Float64} + @test (@inferred x[1:length(x)]) isa Vector{Float64} + @test (@inferred x[[1, 4]]) isa Vector{Float64} - # `similar` with a non-matching dims tuple gives the backing-array fallback + # `similar` with a non-matching dims falls back to the backing array; + # with matching dims keeps the NamedArrayPartition wrapper. @test similar(x, Float64, (2,)) isa Vector{Float64} - @test similar(x, (2,)) isa Vector{Float64} - # `similar` with matching dims preserves the NamedArrayPartition wrapper + @test similar(x, (2,)) isa Vector{Float64} @test similar(x, Float64, size(x)) isa NamedArrayPartition - @test similar(x, size(x)) isa NamedArrayPartition + @test similar(x, size(x)) isa NamedArrayPartition - # Scalar indexing untouched + # Scalar indexing untouched and type-stable @test x[1] == 1.0 @test x[3] == 2.0 + @test (@inferred x[1]) === 1.0 x[1] = 99.0 @test x[1] == 99.0 end