From 8d5548533ccdf7cf773d7164eb8384067267fb99 Mon Sep 17 00:00:00 2001
From: Simeon David Schaub <simeon@schaub.rocks>
Date: Wed, 29 Jul 2026 08:56:37 +0000
Subject: [PATCH] don't return `ReshapedArray` from `mapreducedim!`

`mapreducedim!` reshapes the output container to match the number of
dimensions of the input, but then returned that reshaped array instead of
the one it was handed. Keep a reference to the original and return that.

Same as https://github.com/JuliaGPU/CUDA.jl/pull/3219, discovered in
https://github.com/JuliaGPU/GPUArrays.jl/pull/754.
---
 src/mapreduce.jl | 3 ++-
 test/array.jl    | 6 ++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index 645db2cd..6813fdca 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -117,6 +117,7 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::oneWrappedArray{T},
     Base.check_reducedims(R, A)
     length(A) == 0 && return R # isempty(::Broadcasted) iterates
 
+    R_old = R
     # add singleton dimensions to the output container, if needed
     if ndims(R) < ndims(A)
         dims = Base.fill_to_length(size(R), 1, Val(ndims(A)))
@@ -199,5 +200,5 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::oneWrappedArray{T},
         GPUArrays.mapreducedim!(identity, op, R′, partial; init=init)
     end
 
-    return R
+    return R_old
 end
diff --git a/test/array.jl b/test/array.jl
index 1ca66453..53eaba96 100644
--- a/test/array.jl
+++ b/test/array.jl
@@ -139,3 +139,9 @@ end
   resize!(b, 1)
   @test length(b) == 1
 end
+
+@testset "mapreducedim! returning same type" begin
+  R = transpose(oneAPI.zeros(Float32, 2, 3))
+  A = oneArray(rand(Float32, 3, 2, 10))
+  @test @inferred(oneAPI.GPUArrays.mapreducedim!(identity, +, R, A)) === R
+end