[JLArrays] allow mapreducedim! into wrapper

lib/JLArrays/Project.toml

@@ -1,7 +1,7 @@
 name = "JLArrays"
 uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/JLArrays/src/JLArrays.jl

@@ -555,12 +555,15 @@ else
     end
 end
 
+struct ArrayNoCopy end
+Adapt.adapt_storage(::ArrayNoCopy, x::JLArray) = typed_data(x)
+
 function GPUArrays.mapreducedim!(f, op, R::AnyJLArray, A::Union{AbstractArray,Broadcast.Broadcasted};
                                  init=nothing)
     if init !== nothing
         fill!(R, init)
     end
-    @allowscalar Base.reducedim!(op, typed_data(R), map(f, A))
+    @allowscalar Base.reducedim!(op, adapt(ArrayNoCopy(), R), map(f, A))
     R
 end
 

test/testsuite/reductions.jl

@@ -34,6 +34,18 @@ end
         # implicit singleton dimensions
         @test compare((A,R)->Base.mapreducedim!(identity, +, R, A), AT, rand(range, (2,2)), zeros(ET, (2,)))
         @test compare((A,R)->Base.mapreducedim!(identity, +, R, A), AT, rand(range, (2,3)), zeros(ET, (2,)))
+
+        # mapreducedim! into wrapper types
+        for t in [transpose, adjoint]
+            @test compare((A,R)->Base.mapreducedim!(identity, +, R, A), AT, rand(range, (2,2)), t(zeros(ET, (2,))))
+            @test compare((A,R)->Base.mapreducedim!(abs2, +, R, A), AT, rand(range, (3,2,10)), t(zeros(ET, (2,3))))
+
+            # TODO: reenable once https://github.com/JuliaGPU/Metal.jl/issues/907 is fixed
+            # @test compare((A,R)->sum!(abs2, R, A), AT, rand(range, (3,2,10)), t(zeros(ET, (2,3))))
+
+            A, R = AT(rand(range, (3,2,10))), t(AT(zeros(ET, (2,3))))
+            @test Base.mapreducedim!(identity, *, R, A) === R
+        end
     end
 end