Replace LRP gradient computation with VJP using Zygote.pullback (#72)

Author: adrhill

src/lrp_rules.jl

@@ -11,14 +11,8 @@ function lrp_autodiff!(
     Rₖ::T1, rule::R, layer::L, aₖ::T1, Rₖ₊₁::T2
 ) where {R<:AbstractLRPRule,L,T1,T2}
     layerᵨ = modify_layer(rule, layer)
-    c::T1 = only(
-        gradient(aₖ) do a
-            z::T2 = layerᵨ(a)
-            s = Zygote.@ignore Rₖ₊₁ ./ modify_denominator(rule, z)
-            z ⋅ s
-        end,
-    )
-    Rₖ .= aₖ .* c
+    ãₖ₊₁, back = Zygote.pullback(layerᵨ, aₖ)
+    Rₖ .=  aₖ .* only(back(Rₖ₊₁ ./ modify_denominator(rule, ãₖ₊₁)))
     return nothing
 end
 

test/test_heatmaps.jl

@@ -22,7 +22,7 @@ for r in reducers
 end
 
 @test_throws ArgumentError heatmap(A, reduce=:foo)
-@test_throws ErrorException heatmap(A, rangescale=:bar)
+@test_throws ArgumentError heatmap(A, rangescale=:bar)
 
 B = reshape(A, 2, 2, 3, 1, 1)
 @test_throws DomainError heatmap(B)

test/test_rules.jl

@@ -141,7 +141,7 @@ layers = Dict(
     "MeanPool" => MaxPool((3, 3)),
     "ConvTranspose" => ConvTranspose((3, 3), 2 => 4; init=pseudorandn),
     "CrossCor" => CrossCor((3, 3), 2 => 4; init=pseudorandn),
-    "flatten" => flatten,
+    "flatten" => Flux.flatten,
     "Dropout" => Dropout(0.2),
     "AlphaDropout" => AlphaDropout(0.2),
 )
@@ -178,7 +178,7 @@ lrp!(Rₖ, rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp!(Rₖ, rule,
 layers = Dict(
     "Conv" => (Conv((3, 3), 2 => 4; init=pseudorandn), aₖ),
     "Dense_relu" => (Dense(ins_dense, outs_dense, relu; init=pseudorandn), aₖ_dense),
-    "flatten" => (flatten, aₖ),
+    "flatten" => (Flux.flatten, aₖ),
 )
 @testset "Custom layers" begin
     for (rulename, rule) in RULES