Skip to content

fix: mapreduce type stability, VoA broadcast adjoints#325

Merged
ChrisRackauckas merged 3 commits into
SciML:masterfrom
AayushSabharwal:as/inference
Jan 5, 2024
Merged

fix: mapreduce type stability, VoA broadcast adjoints#325
ChrisRackauckas merged 3 commits into
SciML:masterfrom
AayushSabharwal:as/inference

Conversation

@AayushSabharwal

Copy link
Copy Markdown
Member

Close #323

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

Additional context

Add any other context about the problem here.

Comment thread src/vector_of_array.jl Outdated
@inline Base.sum(VA::AbstractVectorOfArray; kwargs...) = sum(identity, VA; kwargs...)
@inline function Base.sum(f, VA::AbstractVectorOfArray; kwargs...)
if hasproperty(kwargs, :dims)
sum(Array(VA); kwargs...)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This can be a potentially big allocation

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

That was the default implementation, but I just realized two things:

  • that method is probably not necessary
  • mapreduce is implemented wrong

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in bdb5830

Comment thread src/vector_of_array.jl Outdated

Base.map(f, A::RecursiveArrayTools.AbstractVectorOfArray) = map(f, A.u)

function Base.mapreduce(f, op, A::AbstractVectorOfArray)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

probably need to have an init

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in bdb5830

@AayushSabharwal AayushSabharwal changed the title refactor: make mapreduce for ArrayPartition, VoA more type-stable refactor: mapreduce type stability, VoA broadcast adjoints Jan 5, 2024
@AayushSabharwal AayushSabharwal changed the title refactor: mapreduce type stability, VoA broadcast adjoints fix: mapreduce type stability, VoA broadcast adjoints Jan 5, 2024
_minus(Δ) = .-Δ
_minus(::Nothing) = nothing

@adjoint function Broadcast.broadcasted(::typeof(-), x::AbstractVectorOfArray, y::Union{AbstractVectorOfArray, Zygote.Numeric})

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

you're never supposed to define specific broadcasts like this, and I don't think digging into Zygote's broadcast system is the answer. It just relies on the Julia-level broadcast

@ChrisRackauckas

Copy link
Copy Markdown
Member

This will need a real fix, but let's at least get downstream tests passing.

@ChrisRackauckas ChrisRackauckas merged commit ed2ea13 into SciML:master Jan 5, 2024
@AayushSabharwal AayushSabharwal deleted the as/inference branch January 5, 2024 20:29
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.

Type instability when iterating nested recursive arrays

2 participants