From 1527f2de1636b5b1e9f3fd007aa1fca185b006d8 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Tue, 31 Mar 2020 16:58:44 +0200
Subject: [PATCH] Test for implicit singleton dims with mapreducedim.

---
 Manifest.toml    | 4 ++--
 src/mapreduce.jl | 5 +++++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index 3f88c105..d960aa6a 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -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"
diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index 803093bc..14bcfe12 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -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)