Add GPU tests

Author: adrhill

test/Project.toml

@@ -4,9 +4,11 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 PkgJogger = "10150987-6cc1-4b76-abee-b1c1cbd91c01"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"

test/runtests.jl

@@ -33,6 +33,9 @@ using JET
         @info "Testing analyzers on batches..."
         include("test_batches.jl")
     end
+    @testset "GPU tests" begin
+        include("test_gpu.jl")
+    end
     @testset "Benchmark correctness" begin
         @info "Testing whether benchmarks are up-to-date..."
         include("test_benchmarks.jl")

test/test_gpu.jl

@@ -0,0 +1,39 @@
+using ExplainableAI
+using Test
+
+using Flux
+using Metal, JLArrays
+
+if Metal.functional()
+    @info "Using Metal as GPU device"
+    device = mtl # use Apple Metal locally
+else
+    @info "Using JLArrays as GPU device"
+    device = jl # use JLArrays to fake GPU array
+end
+
+model = Chain(Dense(10 => 32, relu), Dense(32 => 5))
+input = rand(Float32, 10, 8)
+@test_nowarn model(input)
+
+model_gpu = device(model)
+input_gpu = device(input)
+@test_nowarn model_gpu(input_gpu)
+
+analyzer_types = (Gradient, SmoothGrad, InputTimesGradient)
+
+@testset "Run analyzer (CPU)" begin
+    @testset "$A" for A in analyzer_types
+        analyzer = A(model)
+        expl = analyze(input, analyzer)
+        @test expl isa Explanation
+    end
+end
+
+@testset "Run analyzer (GPU)" begin
+    @testset "$A" for A in analyzer_types
+        analyzer_gpu = A(model_gpu)
+        expl = analyze(input_gpu, analyzer_gpu)
+        @test expl isa Explanation
+    end
+end