Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix multi-dimensional VectorOfArray broadcast #374

Merged
merged 12 commits into from
May 6, 2024

Conversation

jlchan
Copy link
Contributor

@jlchan jlchan commented May 6, 2024

Addresses #373

The fix just involves modifying Base.copy for VectorOfArray with multi-dimensional parent arrays. This could replace the old Base.copy implementation, but I left the old implementation there just in case it changes behavior upstream.

jlchan added 12 commits May 5, 2024 15:50
the previous implementation for VoA with Vector parent arrays was `VectorOfArray([similar(VA[:, i], T) for i in eachindex(VA.u)])`. This creates a VoA by calling `similar` on each entry of `VA`, but necessitated a different implementation for VoA with multi-dimensional parent arrays. 

Creating a new VoA by broadcasting `similar` over the parent array is (I believe) equivalent to previous implementations for VoA with both Vector and Array parents.
the previous implementation for VoA with Vector parent arrays was `VectorOfArray([similar(VA[:, i], T) for i in eachindex(VA.u)])`. This creates a VoA by calling `similar` on each entry of `VA`, but necessitated a different implementation for VoA with multi-dimensional parent arrays. 

Creating a new VoA by broadcasting `similar` over the parent array is (I believe) equivalent to previous implementations for VoA with both Vector and Array parents.
This reverts commit f4675f2.
Comment on lines +863 to +873
if parent isa AbstractVector
# this is the default behavior in v3.15.0
N = narrays(bc)
return VectorOfArray(map(1:N) do i
copy(unpack_voa(bc, i))
end)
else # if parent isa AbstractArray
return VectorOfArray(map(enumerate(Iterators.product(axes(parent)...))) do (i, _)
copy(unpack_voa(bc, i))
end)
end
Copy link
Contributor Author

@jlchan jlchan May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Side note, this entire block can be replaced by the second branch

return VectorOfArray(map(enumerate(Iterators.product(axes(parent)...))) do (i, _)
    copy(unpack_voa(bc, i))
end)

since enumerate(Iterators.product(axes(parent)...))) do (i, _) basically recovers map(1:N) do i for an AbstractVector.

I left the old behavior in since I wasn't sure if it would break something upstream or change performance (and it's easier to read).

@@ -719,7 +718,7 @@ end
# for VectorOfArray with multi-dimensional parent arrays of arrays where all elements are the same type
function Base.similar(vec::VectorOfArray{
T, N, AT}) where {T, N, AT <: AbstractArray{<:AbstractArray{T}}}
return VectorOfArray(similar(Base.parent(vec)))
return VectorOfArray(similar.(Base.parent(vec)))
Copy link
Contributor Author

@jlchan jlchan May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Broadcasting similar here avoids some #undef issues.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Side note: this implementation could replace these 3 functions

@inline function Base.similar(VA::VectorOfArray, ::Type{T} = eltype(VA)) where {T}
VectorOfArray([similar(VA[:, i], T) for i in eachindex(VA.u)])
end
# for VectorOfArray with multi-dimensional parent arrays of arrays where all elements are the same type
function Base.similar(vec::VectorOfArray{
T, N, AT}) where {T, N, AT <: AbstractArray{<:AbstractArray{T}}}
return VectorOfArray(similar(Base.parent(vec)))
end
# special-case when the multi-dimensional parent array is just an AbstractVector (call the old method)
function Base.similar(vec::VectorOfArray{
T, N, AT}) where {T, N, AT <: AbstractVector{<:AbstractArray{T}}}
return Base.similar(vec, eltype(vec))
end

However, doing so breaks these tests because similar can't be called on <:Real values:

x = [1, 2, 3, 4, 5, 6, 7, 8, 9]
testva = DiffEqArray(x, x)
testvb = DiffEqArray(x, x)
mulX = sqrt.(abs.(testva .* testvb))
ref = sqrt.(abs.(x .* x))
@test mulX == ref
fill!(mulX, 0)
mulX .= sqrt.(abs.(testva .* testvb))
@test mulX == ref

These tests don't seem correct to me; the parent array x doesn't have underlying data structure Vector{AbstractArray{T}}, so it shouldn't be a valid edge case. Maybe I'm missing something?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Numbers are also allowed. That's used a lot downstream.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can special-case that case where the elements are numbers, but we cannot throw it out.

@jlchan
Copy link
Contributor Author

jlchan commented May 6, 2024

Applying formatting introduces a bunch of spurious file changes so I'm skipping it for now.

@jlchan jlchan marked this pull request as ready for review May 6, 2024 05:53
@ChrisRackauckas ChrisRackauckas merged commit f7415e8 into SciML:master May 6, 2024
18 of 25 checks passed
@jlchan
Copy link
Contributor Author

jlchan commented May 6, 2024

Thanks for the quick review!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants