Test Tracker.jacobian on GPU

FluxML · May 30, 2019 · 8255327 · 8255327
1 parent 8ee6af1
commit 8255327
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/test/cuda/cuda.jl b/test/cuda/cuda.jl
@@ -48,6 +48,13 @@ end
   @test y[3,:] isa CuArray
 end
 
+@testset "Jacobian on GPU" begin
+  # https://github.com/FluxML/Tracker.jl/pull/33
+  @test collect(jacobian(identity, gpu([0.0, 0.0]))) == [1 0; 0 1]
+  @test collect(gradient(x -> sum(jacobian(y -> y .^ 2, x) .^ 2),
+                         gpu([1.0, 2.0, 3.0]))) == [8.0, 16.0, 24.0]
+end
+
 if CuArrays.libcudnn != nothing
     @info "Testing Flux/CUDNN"
     include("cudnn.jl")