Skip to content
This repository was archived by the owner on Mar 12, 2021. It is now read-only.

Test for implicit singleton dims with mapreducedim. #665

Merged
merged 1 commit into from
Mar 31, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Manifest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,8 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[GPUArrays]]
deps = ["AbstractFFTs", "Adapt", "LinearAlgebra", "Printf", "Random", "Serialization"]
git-tree-sha1 = "fb4ae63d16ba6a8324a0a01778c2cbea2b25a510"
repo-rev = "88114d40ed5a594f6af1939a89bbbf59c79c6fd4"
git-tree-sha1 = "50542dca6e8339a5e0a6718283f956187123234a"
repo-rev = "cb79e08c09ca0eb776c1ded7b7fe8876bd012981"
repo-url = "https://github.com/JuliaGPU/GPUArrays.jl.git"
uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
version = "3.1.0"
Expand Down
5 changes: 5 additions & 0 deletions src/mapreduce.jl
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,11 @@ NVTX.@range function GPUArrays.mapreducedim!(f, op, R::CuArray{T}, As::AbstractA
shuffle &= capability(device()) >= v"3.0"
shuffle &= T in (Bool, Int32, Int64, Float32, Float64, ComplexF32, ComplexF64)

# add singleton dimensions to the output container, if needed
if ndims(R) < ndims(A)
R = reshape(R, ntuple(i -> ifelse(i <= ndims(R), size(R,i), 1), ndims(A)))
end

# iteration domain, split in two: one part covers the dimensions that should
# be reduced, and the other covers the rest. combining both covers all values.
Rall = CartesianIndices(A)
Expand Down