diff --git a/src/host/broadcast.jl b/src/host/broadcast.jl index 7132863d..1c7dfc08 100644 --- a/src/host/broadcast.jl +++ b/src/host/broadcast.jl @@ -109,14 +109,14 @@ function Base.map(f, x1::AnyGPUArray, xrest::Union{AnyGPUArray, AbstractRange}.. return map!(f, dest, xs...) end -function Base.map!(f, dest::AnyGPUArray, xs::AbstractArray...) +function Base.map!(f, dest::AnyGPUArray, x::AbstractArray, xs::AbstractArray...) # custom broadcast, ignoring the container size mismatches # (avoids the reshape + view that our mapreduce impl has to do) - indices = LinearIndices.((dest, xs...)) + indices = LinearIndices.((dest, x, xs...)) common_length = minimum(length.(indices)) common_length==0 && return - bc = Broadcast.instantiate(Broadcast.broadcasted(f, xs...)) + bc = Broadcast.instantiate(Broadcast.broadcasted(f, x, xs...)) if bc isa Broadcast.Broadcasted bc = Broadcast.preprocess(dest, bc) end diff --git a/test/testsuite/base.jl b/test/testsuite/base.jl index 77e92135..42d22bb1 100644 --- a/test/testsuite/base.jl +++ b/test/testsuite/base.jl @@ -285,6 +285,11 @@ end @test compare((a, b)-> map(+, a, b), AT, rand(Float32, 10), rand(Float32, 10)) @test compare((a, b)-> map!(-, a, b), AT, rand(Float32, 10), rand(Float32, 10)) @test compare((a, b, c, d)-> map!(*, a, b, c, d), AT, rand(Float32, 10), rand(Float32, 10), rand(Float32, 10), rand(Float32, 10)) + if VERSION >= v"1.12" + @test compare(a-> map!(sqrt, a), AT, rand(Float32, 10)) + else + @test_throws ArgumentError map!(sqrt, AT(rand(Float32, 10))) + end end @testset "repeat" begin