Test type stability with JET.jl (#163)

* Add JET and JuliaFormatter tests

* Make individual tests modular for use with TestEnv.jl

Author: adrhill

Project.toml

@@ -13,7 +13,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Distributions = "0.25"
-Random = "1"
+Random = "<0.0.1, 1"
 Reexport = "1"
 Statistics = "1"
 XAIBase = "3"

README.md

@@ -3,7 +3,7 @@ ___
 
 | **Documentation** | **Build Status** | **DOI** |
 |:----------------- |:---------------- |:------- |
-| [![][docs-stab-img]][docs-stab-url] [![][docs-dev-img]][docs-dev-url] | [![][ci-img]][ci-url] [![][codecov-img]][codecov-url] [![Aqua QA][aqua-img]][aqua-url] | [![][doi-img]][doi-url] |
+| [![][docs-stab-img]][docs-stab-url] [![][docs-dev-img]][docs-dev-url] | [![][ci-img]][ci-url] [![][codecov-img]][codecov-url] [![Aqua][aqua-img]][aqua-url] [![JET][jet-img]][jet-url] | [![][doi-img]][doi-url] |
 
 Explainable AI in Julia.
 
@@ -165,6 +165,9 @@ Contributions are welcome!
 [aqua-img]: https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg
 [aqua-url]: https://github.com/JuliaTesting/Aqua.jl
 
+[jet-img]: https://img.shields.io/badge/%F0%9F%9B%A9%EF%B8%8F_tested_with-JET.jl-233f9a
+[jet-url]: https://github.com/aviatesk/JET.jl
+
 [docs-composites]: https://julia-xai.github.io/ExplainableAI.jl/stable/generated/lrp/composites/
 [docs-custom-rules]: https://julia-xai.github.io/ExplainableAI.jl/stable/generated/lrp/custom_rules/
 

test/Project.toml

@@ -2,11 +2,12 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 XAIBase = "9b48221d-a747-4c1b-9860-46a1d8ba24a7"
-

test/runtests.jl

@@ -1,17 +1,27 @@
 using ExplainableAI
-using Flux
+
 using Test
-using ReferenceTests
+using JuliaFormatter
 using Aqua
-using Random
-
-pseudorand(dims...) = rand(MersenneTwister(123), Float32, dims...)
+using JET
 
 @testset "ExplainableAI.jl" begin
-    @testset "Aqua.jl" begin
-        @info "Running Aqua.jl's auto quality assurance tests. These might print warnings from dependencies."
-        Aqua.test_all(ExplainableAI; ambiguities=false)
+    @info "Testing formalities..."
+    if VERSION >= v"1.10"
+        @testset "Code formatting" begin
+            @info "- Testing code formatting with JuliaFormatter..."
+            @test JuliaFormatter.format(ExplainableAI; verbose=false, overwrite=false)
+        end
+        @testset "Aqua.jl" begin
+            @info "- Running Aqua.jl tests. These might print warnings from dependencies..."
+            Aqua.test_all(ExplainableAI; ambiguities=false)
+        end
+        @testset "JET tests" begin
+            @info "- Testing type stability with JET..."
+            JET.test_package(ExplainableAI; target_defined_modules=true)
+        end
     end
+
     @testset "Input augmentation" begin
         @info "Testing input augmentation..."
         include("test_input_augmentation.jl")

test/test_batches.jl

@@ -1,8 +1,12 @@
-using Flux
 using ExplainableAI
+using Test
+
+using Flux
 using Random
 using Distributions: Laplace
 
+pseudorand(dims...) = rand(MersenneTwister(123), Float32, dims...)
+
 ## Test `fuse_batchnorm` on Dense and Conv layers
 ins = 20
 outs = 10

test/test_cnn.jl

@@ -1,5 +1,9 @@
 using ExplainableAI
+using Test
+using ReferenceTests
+
 using Flux
+using Random
 using JLD2
 
 const GRADIENT_ANALYZERS = Dict(
@@ -9,6 +13,8 @@ const GRADIENT_ANALYZERS = Dict(
     "GradCAM"             => m -> GradCAM(m[1], m[2]),
 )
 
+pseudorand(dims...) = rand(MersenneTwister(123), Float32, dims...)
+
 input_size = (32, 32, 3, 1)
 input = pseudorand(input_size)
 

test/test_input_augmentation.jl

@@ -1,5 +1,6 @@
 using ExplainableAI: augment_batch_dim, augment_indices, reduce_augmentation
 using ExplainableAI: interpolate_batch
+using Test
 
 # augment_batch_dim
 A = [1 2; 3 4]