add some GPU tests

Author: mcabbott

test/rulesets/Base/broadcast.jl

@@ -80,21 +80,21 @@ BT1 = Broadcast.BroadcastStyle(Tuple)
 
     @testset "fused rules" begin
         @testset "arithmetic" begin
-            test_rrule(copy∘broadcasted, +, rand(3), rand(3))
-            test_rrule(copy∘broadcasted, +, rand(3), rand(4)')
-            test_rrule(copy∘broadcasted, +, rand(3), rand(1), rand())
-            test_rrule(copy∘broadcasted, +, rand(3), 1.0*im)
-            test_rrule(copy∘broadcasted, +, rand(3), true)
-            test_rrule(copy∘broadcasted, +, rand(3), Tuple(rand(3)))
+            @gpu test_rrule(copy∘broadcasted, +, rand(3), rand(3))
+            @gpu test_rrule(copy∘broadcasted, +, rand(3), rand(4)')
+            @gpu test_rrule(copy∘broadcasted, +, rand(3), rand(1), rand())
+            @gpu test_rrule(copy∘broadcasted, +, rand(3), 1.0*im)
+            @gpu test_rrule(copy∘broadcasted, +, rand(3), true)
+            @gpu_broken test_rrule(copy∘broadcasted, +, rand(3), Tuple(rand(3)))
 
-            test_rrule(copy∘broadcasted, -, rand(3), rand(3))
-            test_rrule(copy∘broadcasted, -, rand(3), rand(4)')
-            test_rrule(copy∘broadcasted, -, rand(3))
+            @gpu test_rrule(copy∘broadcasted, -, rand(3), rand(3))
+            @gpu test_rrule(copy∘broadcasted, -, rand(3), rand(4)')
+            @gpu test_rrule(copy∘broadcasted, -, rand(3))
             test_rrule(copy∘broadcasted, -, Tuple(rand(3)))
 
-            test_rrule(copy∘broadcasted, *, rand(3), rand(3))
-            test_rrule(copy∘broadcasted, *, rand(3), rand())
-            test_rrule(copy∘broadcasted, *, rand(), rand(3))
+            @gpu test_rrule(copy∘broadcasted, *, rand(3), rand(3))
+            @gpu test_rrule(copy∘broadcasted, *, rand(3), rand())
+            @gpu test_rrule(copy∘broadcasted, *, rand(), rand(3))
 
             test_rrule(copy∘broadcasted, *, rand(3) .+ im, rand(3) .+ 2im)
             test_rrule(copy∘broadcasted, *, rand(3) .+ im, rand() + 3im)
@@ -107,14 +107,15 @@ BT1 = Broadcast.BroadcastStyle(Tuple)
             @test unthunk(bk4([4, 5im, 6+7im])[4]) == [0,5,7]
 
             # These two test vararg rrule * rule:
-            test_rrule(copy∘broadcasted, *, rand(3), rand(3), rand(3), rand(3), rand(3))
-            test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)')
+            @gpu test_rrule(copy∘broadcasted, *, rand(3), rand(3), rand(3), rand(3), rand(3))
+            @gpu_broken test_rrule(copy∘broadcasted, *, rand(), rand(), rand(3), rand(3) .+ im, rand(4)')
+            # GPU error from dot(x::JLArray{Float32, 1}, y::JLArray{ComplexF32, 2})
 
-            test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3), Val(2))
-            test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2))
+            @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3), Val(2))
+            @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2))
 
-            test_rrule(copy∘broadcasted, /, rand(3), rand())
-            test_rrule(copy∘broadcasted, /, rand(3) .+ im, rand() + 3im)
+            @gpu test_rrule(copy∘broadcasted, /, rand(3), rand())
+            @gpu test_rrule(copy∘broadcasted, /, rand(3) .+ im, rand() + 3im)
         end
         @testset "identity etc" begin
             test_rrule(copy∘broadcasted, identity, rand(3))

test/unzipped.jl

@@ -79,4 +79,19 @@ using ChainRules: unzip_broadcast, unzip #, unzip_map
         @test dx5[2] isa Tangent{<:Tuple}
         @test Tuple(dx5[2][2]) == (10, ZeroTangent())
     end
+    
+    @testset "JLArray tests" begin  # fake GPU testing
+        (y1, y2), bk = rrule(CFG, unzip_broadcast, tuple, [1,2,3.0], [4 5.0])
+        (y1jl, y2jl), bk_jl = rrule(CFG, unzip_broadcast, tuple, jl([1,2,3.0]), jl([4 5.0]))
+        @test y1 == Array(y1jl)
+        # TODO invent some tests of this rrule's pullback function
+
+        @test unzip(jl([(1,2), (3,4), (5,6)])) == (jl([1, 3, 5]), jl([2, 4, 6]))
+
+        @test unzip(jl([(missing,2), (missing,4), (missing,6)]))[2] == jl([2, 4, 6])
+        @test unzip(jl([(missing,2), (missing,4), (missing,6)]))[2] isa Base.ReinterpretArray
+
+        @test unzip(jl([(1,), (3,), (5,)]))[1] == jl([1, 3, 5])
+        @test unzip(jl([(1,), (3,), (5,)]))[1] isa Base.ReinterpretArray
+    end
 end