Faster stabilize_denom (#47)

Author: adrhill

src/lrp_rules.jl

@@ -54,7 +54,7 @@ modify_params(::AbstractLRPRule, W, b) = (W, b) # general fallback
 
 Function that modifies zₖ on the forward pass, e.g. for numerical stability.
 """
-modify_denominator(::AbstractLRPRule, d) = stabilize_denom(d; eps=1.0f-9) # general fallback
+modify_denominator(::AbstractLRPRule, d) = stabilize_denom(d, 1.0f-9) # general fallback
 
 """
     modify_layer(rule, layer)
@@ -104,7 +104,7 @@ struct EpsilonRule{T} <: AbstractLRPRule
     ϵ::T
     EpsilonRule(; ϵ=1.0f-6) = new{Float32}(ϵ)
 end
-modify_denominator(r::EpsilonRule, d) = stabilize_denom(d; eps=r.ϵ)
+modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
 
 """
     ZBoxRule()
@@ -131,7 +131,7 @@ function lrp_zbox!(layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
         f⁻::T2 = layer⁻(h)
 
         z = f - f⁺ - f⁻
-        s = Zygote.@ignore safedivide(Rₖ₊₁, z; eps=1e-9)
+        s = Zygote.@ignore safedivide(Rₖ₊₁, z, 1e-9)
         z ⋅ s
     end
     Rₖ .= aₖ .* c + l .* cₗ + h .* cₕ

src/utils.jl

@@ -1,20 +1,23 @@
 """
-    stabilize_denom(d; eps = 1f-6)
+    stabilize_denom(d, [eps = 1f-9])
 
 Replace zero terms of a matrix `d` with `eps`.
 """
-stabilize_denom(d::Real; eps=1.0f-9) = ifelse(d ≈ 0, d + sign(d) * eps, d)
-function stabilize_denom(d::AbstractArray; eps=1.0f-9)
-    return d + ((d .≈ 0) + sign.(d)) * eps
+function stabilize_denom(d::T, eps=T(1.0f-9)) where {T}
+    iszero(d) && (return T(eps))
+    return d + sign(d) * T(eps)
 end
+stabilize_denom(D::AbstractArray{T}, eps=T(1.0f-9)) where {T} = stabilize_denom.(D, eps)
 
 """
-    safedivide(a, b; eps = 1f-6)
+    safedivide(a, b, [eps = 1f-6])
 
 Elementwise division of two matrices avoiding near zero terms
 in the denominator by replacing them with `± eps`.
 """
-safedivide(a, b; eps=1.0f-9) = a ./ stabilize_denom(b; eps=eps)
+function safedivide(a::AbstractArray{T}, b::AbstractArray{T}, eps=T(1.0f-9)) where {T}
+    return a ./ stabilize_denom(b, T(eps))
+end
 
 """
     batch_dim_view(A)

test/test_utils.jl

@@ -1,6 +1,6 @@
 using Flux
-using ExplainableAI: flatten_model, has_output_softmax, check_output_softmax
-using ExplainableAI: stabilize_denom
+using ExplainableAI:
+    flatten_model, has_output_softmax, check_output_softmax, stabilize_denom
 
 # flatten_model
 @test flatten_model(Chain(Chain(Chain(abs)), sqrt, Chain(relu))) == Chain(abs, sqrt, relu)
@@ -22,5 +22,8 @@ using ExplainableAI: stabilize_denom
     Chain(Chain(abs), Chain(Chain(softmax)), sqrt) # don't do anything if there is no softmax at the end
 
 # stabilize_denom
-A = [1.0 0; -0 -1.0e-25]
-@test stabilize_denom(A; eps=1e-3) ≈ [1.001 1e-3; 1e-3 -1e-3]
+A = [1.0 0.0 1.0e-25; -1.0 -0.0 -1.0e-25]
+S = @inferred stabilize_denom(A, 1e-3)
+@test S ≈ [1.001 1e-3 1e-3; -1.001 1e-3 -1e-3]
+S = @inferred stabilize_denom(Float32.(A), 1e-2)
+@test S ≈ [1.01 1f-2 1f-2; -1.01 1f-2 -1f-2]