-
Notifications
You must be signed in to change notification settings - Fork 148
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
base: master
Are you sure you want to change the base?
Conversation
@inline _mapreduce(f, op, D::Tuple{<:Any}, init, sz::Size{S}, a::StaticArray) where {S} = | ||
_mapreduce(f, op, first(D), init, sz, a) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
Closes #880.