Make Tullio weak dependency using package extensions (#141)

* Make Tullio weak dependency using package extensions. For backwards compatibility, it is still a normal dependency on Julia versions prior to 1.9.

* Skip Aqua.jl's Project.toml tests on Julia 1.6

Author: adrhill

Project.toml

@@ -16,6 +16,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
+
+[extensions]
+TullioLRPRulesExt = "Tullio"
+
 [compat]
 ColorSchemes = "3.18"
 Distributions = "0.25"

ext/TullioLRPRulesExt.jl

@@ -0,0 +1,21 @@
+module TullioLRPRulesExt
+
+using ExplainableAI, Flux, Tullio
+import ExplainableAI: lrp!, modify_input, modify_denominator
+import ExplainableAI: ZeroRule, EpsilonRule, GammaRule, WSquareRule
+
+# Fast implementation for Dense layer using Tullio.jl's einsum notation:
+for R in (ZeroRule, EpsilonRule, GammaRule)
+    @eval function lrp!(Rᵏ, rule::$R, layer::Dense, modified_layer, aᵏ, Rᵏ⁺¹)
+        layer = isnothing(modified_layer) ? layer : modified_layer
+        ãᵏ = modify_input(rule, aᵏ)
+        z = modify_denominator(rule, layer(ãᵏ))
+        @tullio Rᵏ[j, b] = layer.weight[i, j] * ãᵏ[j, b] / z[i, b] * Rᵏ⁺¹[i, b]
+    end
+end
+
+function lrp!(Rᵏ, ::WSquareRule, _layer::Dense, modified_layer::Dense, aᵏ, Rᵏ⁺¹)
+    den = sum(modified_layer.weight; dims=2)
+    @tullio Rᵏ[j, b] = modified_layer.weight[i, j] / den[i] * Rᵏ⁺¹[i, b]
+end
+end # module

src/ExplainableAI.jl

@@ -6,7 +6,6 @@ using Distributions: Distribution, Sampleable, Normal
 using Random: AbstractRNG, GLOBAL_RNG
 using Flux
 using Zygote
-using Tullio
 using Markdown
 
 # Heatmapping:
@@ -68,4 +67,12 @@ export heatmap
 # utils
 export strip_softmax, flatten_model, canonize
 export preprocess_imagenet
+
+# Package extension backwards compatibility with Julia 1.6.
+# For Julia 1.6, Tullio is treated as a normal dependency and always loaded.
+# https://pkgdocs.julialang.org/v1/creating-packages/#Transition-from-normal-dependency-to-extension
+if !isdefined(Base, :get_extension)
+    include("../ext/TullioLRPRulesExt.jl")
+end
+
 end # module

src/lrp/rules.jl

@@ -540,17 +540,5 @@ function lrp!(Rᵏ, _rule::FlatRule, _layer::Dense, _modified_layer, _aᵏ, Rᵏ
     end
 end
 
-# Fast implementation for Dense layer using Tullio.jl's einsum notation:
-for R in (ZeroRule, EpsilonRule, GammaRule)
-    @eval function lrp!(Rᵏ, rule::$R, layer::Dense, modified_layer, aᵏ, Rᵏ⁺¹)
-        layer = isnothing(modified_layer) ? layer : modified_layer
-        ãᵏ = modify_input(rule, aᵏ)
-        z = modify_denominator(rule, layer(ãᵏ))
-        @tullio Rᵏ[j, b] = layer.weight[i, j] * ãᵏ[j, b] / z[i, b] * Rᵏ⁺¹[i, b]
-    end
-end
-
-function lrp!(Rᵏ, ::WSquareRule, _layer::Dense, modified_layer::Dense, aᵏ, Rᵏ⁺¹)
-    den = sum(modified_layer.weight; dims=2)
-    @tullio Rᵏ[j, b] = modified_layer.weight[i, j] / den[i] * Rᵏ⁺¹[i, b]
-end
+# Fast implementations for Dense layers can be conditionally loaded with Tullio.jl
+# using package extensions, see exp/TullioRulesExt.jl

test/Project.toml

@@ -11,13 +11,16 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
+Aqua = "0.7"
 Distributions = "0.25"
 Flux = "0.13, 0.14"
 ImageCore = "0.9, 0.10"
 JLD2 = "0.4"
 LoopVectorization = "0.12"
 Metalhead = "0.8"
-ReferenceTests = "0.10.1"
+ReferenceTests = "0.10"
 Suppressor = "0.2"
+Tullio = "0.3"

test/runtests.jl

@@ -9,10 +9,13 @@ using Random
 pseudorand(dims...) = rand(MersenneTwister(123), Float32, dims...)
 
 @testset "ExplainableAI.jl" begin
-    # Run Aqua.jl quality assurance tests
     @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)
+        # Package extensions break Project.toml formatting tests on Julia 1.6
+        # https://github.com/JuliaTesting/Aqua.jl/issues/105
+        Aqua.test_all(
+            ExplainableAI; ambiguities=false, project_toml_formatting=VERSION >= v"1.7"
+        )
     end
 
     # Run package tests

test/test_rules.jl

@@ -238,6 +238,21 @@ layers = Dict(
     end
 end
 
+# Test with loaded package extension
+using LoopVectorization
+using Tullio
+@testset "Dense Tullio" begin
+    for (rulename, rule) in RULES
+        @testset "$rulename" begin
+            for (layername, layer) in layers
+                @testset "$layername" begin
+                    run_rule_tests(rule, layer, rulename, layername, aᵏ_dense)
+                end
+            end
+        end
+    end
+end
+
 ## Test ConvLayers and others
 cin, cout = 3, 4
 insize = (6, 6, 3, batchsize)