Aligning broadcasting errors between GPU and CPU (#471)

Co-authored-by: Valentin Churavy <vchuravy@users.noreply.github.com>
Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: 3 people

src/host/indexing.jl

@@ -68,6 +68,13 @@ function _setindex!(dest::AbstractGPUArray, src, Is...)
     idims = length.(Is)
     len = prod(idims)
     len==0 && return dest
+    if length(src) != len
+        if length(src) == 1
+            throw(ArgumentError("indexed assignment with a single value to possibly many locations is not supported; perhaps use broadcasting `.=` instead?"))
+        else
+            throw(DimensionMismatch("dimensions must match: a has "*string(length(src))*" elements, b has  "*string(len)))
+        end
+    end
 
     AT = typeof(dest).name.wrapper
     # NOTE: we are pretty liberal here supporting non-GPU sources and indices...

test/testsuite/indexing.jl

@@ -118,6 +118,14 @@ end
     @testset "JuliaGPU/CUDA.jl#461: sliced setindex" begin
         @test compare((X,Y)->(X[1,:] = Y), AT, zeros(Float32, 2,2), ones(Float32, 2))
     end
+
+    @testset "Broadcasting exceptions" for T in eltypes
+        x = AT(zeros(T, (10, 10, 10, 10)))
+        @test_throws ArgumentError x[1, :, :, :] = 0
+        y = AT(rand(T, (5, 5, 5, 5)))
+        @test_throws DimensionMismatch x[1:9,1:9,:,:] = y
+    end
+
 end
 
 @testsuite "indexing find" (AT, eltypes)->begin