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

Allow dims::Tuple in sum #909

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Conversation

mcabbott
Copy link
Collaborator

@mcabbott mcabbott commented May 8, 2021

Closes #880.

test/mapreduce.jl Outdated Show resolved Hide resolved
@mcabbott mcabbott marked this pull request as draft May 8, 2021 19:34
@mcabbott mcabbott marked this pull request as ready for review May 9, 2021 15:43
Comment on lines +161 to +162
@inline _mapreduce(f, op, D::Tuple{<:Any}, init, sz::Size{S}, a::StaticArray) where {S} =
_mapreduce(f, op, first(D), init, sz, a)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Things like sum(abs2, SA[1 2; 3 4], dims = (1,2)) still won't work, that seems like a bigger project. It would be easy to add a method here to make them give a friendly "not yet supported" error, instead of an "internal" error. But opinions vary as to whether that's a good idea.

Copy link
Contributor

Choose a reason for hiding this comment

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

The change needed is not large though, we can use Base.mapreducedim! to generated the code we want:

@generated function _mapfoldl(f, op, dims::Val{D}, init,
                               ::Size{S}, a::StaticArray) where {S,D}
    _init = init === _InitialValue ? :default : :init
    iter = CartesianIndices(S)
    exprs = fill!(similar(iter, Any, Base.reduced_indices(iter, D)), _init)
    rf(ex, nex) = ex === :default ? :(Base.reduce_first(op, $nex)) : :(op($ex,$nex))
    Base.mapreducedim!(I->:(f(a[$(I.I...)])), rf, exprs, iter)
    return quote
        @_inline_meta
        @inbounds elements = tuple($(exprs...))
        @inbounds return similar_type(a, eltype(elements), Size($(size(exprs))))(elements)
    end
end

The current test passed locally, and the influence on compile time seems negligible.

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.

Summing multiple dimensions of a Static Array with sum() throws an error
2 participants