Merge pull request #36 from LuxDL/ap/fmode

Fix ForwardMode tests

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.3.4"
+version = "0.3.5"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

ext/LuxLibForwardDiffExt.jl

@@ -62,4 +62,8 @@ for op in [:conv, :depthwiseconv]
     end
 end
 
+function LuxLib._drop_forwarddiff_partials(x::AbstractArray{<:Dual})
+    return ForwardDiff.value.(x)
+end
+
 end

ext/LuxLibLuxCUDAExt.jl

@@ -8,7 +8,6 @@ import LuxLib: batchnorm, _batchnorm_cudnn!, _get_batchnorm_statistics, FP_32_64
 LuxLib._replicate(rng::CUDA.RNG) = deepcopy(rng)
 
 # api/batchnorm.jl
-
 const CUDNN_BN_ARRAY_TYPE = Union{CuArray{<:FP_32_64, 2}, CuArray{<:FP_32_64, 4},
     CuArray{<:FP_32_64, 5}}
 const BNParamType = Union{Nothing, CuVector{<:FP_32_64}}

src/api/batchnorm.jl

@@ -42,8 +42,9 @@ function batchnorm(x::AA{<:Real, N}, scale::NOrAVR, bias::NOrAVR, running_mean::
     running_var::NOrAVR; momentum::Real, training::Val, epsilon::Real) where {N}
     x_, xm, xv = _normalization(x, running_mean, running_var, scale, bias,
         _get_batchnorm_reduce_dims(x), training, momentum, epsilon)
-
-    return x_, (; running_mean=xm, running_var=xv)
+    stats = (; running_mean=_drop_forwarddiff_partials(xm),
+        running_var=_drop_forwarddiff_partials(xv))
+    return (x_, stats)
 end
 
 @generated function _get_batchnorm_reduce_dims(::AA{T, N}) where {T, N}

src/utils.jl

@@ -98,3 +98,13 @@ function Base.showerror(io::IO, ex::OutdatedNNlibDependencyException)
     print(io, "OutdatedNNlibDependencyException: ")
     return println(io, "$msg")
 end
+
+# Droping ForwardDiff Gradients
+function _drop_forwarddiff_partials end
+
+_drop_forwarddiff_partials(x::AbstractArray) = x
+_drop_forwarddiff_partials(::Nothing) = nothing
+_drop_forwarddiff_partials(x::Tuple) = _drop_forwarddiff_partials.(x)
+function _drop_forwarddiff_partials(x::NamedTuple{N}) where {N}
+    return NamedTuple{N}(map(_drop_forwarddiff_partials, values(x)))
+end

test/jvp.jl

@@ -33,16 +33,16 @@ function jvp_zygote(f, x, u)
     return Jₓ * vec(u)
 end
 
-function test_jvp_computation(f, x, u)
+function test_jvp_computation(f, x, u, on_gpu)
     jvp₁ = jvp_forwarddiff(f, x, u)
-    if !(x isa ComponentArray)
+    if !(x isa ComponentArray && on_gpu)
         # ComponentArray + ForwardDiff on GPU don't play nice
         jvp₂ = jvp_forwarddiff_concrete(f, x, u)
         @test check_approx(jvp₁, jvp₂; atol=1e-5, rtol=1e-5)
-    end
 
-    jvp₃ = jvp_zygote(f, x, u)
-    @test check_approx(jvp₁, jvp₃; atol=1e-5, rtol=1e-5)
+        jvp₃ = jvp_zygote(f, x, u)
+        @test check_approx(jvp₁, jvp₃; atol=1e-5, rtol=1e-5)
+    end
 end
 
 @testset "$mode: Jacobian Vector Products" for (mode, aType, on_gpu) in MODES
@@ -66,9 +66,10 @@ end
             uw = randn(Float32, size(w)...) |> aType
             u = randn(Float32, length(x) + length(w)) |> aType
 
-            test_jvp_computation(x -> op(x, w; flipped), x, ux)
-            test_jvp_computation(w -> op(x, w; flipped), w, uw)
-            test_jvp_computation(xw -> op(xw.x, xw.w; flipped), ComponentArray(; x, w), u)
+            test_jvp_computation(x -> op(x, w; flipped), x, ux, on_gpu)
+            test_jvp_computation(w -> op(x, w; flipped), w, uw, on_gpu)
+            test_jvp_computation(xw -> op(xw.x, xw.w; flipped), ComponentArray(; x, w), u,
+                on_gpu)
         end
     end
 end