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