feat: more nested AD rules (#1151)

* feat: softmax and logsoftmax jvp rules * feat: add pooling rules * test: logsoftmax and softmax forwarddiff rules * fix: patch meanpool * test: more tests fixed
LuxDL · Jan 1, 2025 · 63d3434 · 63d3434 · avik-pal · Jan 1, 2025
1 parent 3c3a432
commit 63d3434
Show file tree

Hide file tree

Showing 7 changed files with 158 additions and 29 deletions.
diff --git a/Project.toml b/Project.toml
@@ -89,7 +89,7 @@ ComponentArrays = "0.15.18"
 ConcreteStructs = "0.2.3"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 FastClosures = "0.3.2"
 Flux = "0.15, 0.16"
 ForwardDiff = "0.10.36"

diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.3.11"
+version = "1.4.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -66,7 +66,7 @@ Compat = "4.16"
 CpuId = "0.3"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"

diff --git a/lib/LuxLib/ext/LuxLibCUDAExt/LuxLibCUDAExt.jl b/lib/LuxLib/ext/LuxLibCUDAExt/LuxLibCUDAExt.jl
@@ -1,12 +1,22 @@
 module LuxLibCUDAExt
 
-using CUDA: CUDA, CUBLAS, StridedCuMatrix, StridedCuVector, CuPtr
+using CUDA: CUDA, CUBLAS, CuArray, StridedCuMatrix, StridedCuVector, CuPtr
+using ForwardDiff: ForwardDiff
 using LinearAlgebra: LinearAlgebra, Transpose, Adjoint
-using LuxLib: LuxLib, Optional
+using LuxLib: LuxLib, Impl, Optional
 using LuxLib.Utils: ofeltype_array
 using NNlib: NNlib
 using Static: True, False
 
+# Hacky Type Piracy for ForwardDiff rules
+for op in (:logsoftmax, :softmax)
+    dual_op = Symbol(op, :_dual)
+    @eval function NNlib.$(op)(
+            x::CuArray{<:ForwardDiff.Dual{Tag, T, P}}; dims=1) where {Tag, T, P}
+        return Impl.$(dual_op)(x; dims)
+    end
+end
+
 # Low level functions
 include("cublaslt.jl")
 

diff --git a/lib/LuxLib/src/impl/forward_diff.jl b/lib/LuxLib/src/impl/forward_diff.jl
@@ -1,4 +1,4 @@
-for op in [:conv, :depthwiseconv, :∇conv_data, :∇conv_filter]
+for op in (:conv, :depthwiseconv, :∇conv_data, :∇conv_filter)
     patched_op = op !== :depthwiseconv ? eval(op) : getfield(NNlib, op)
 
     @eval function NNlib.$(op)(x1::AbstractArray{<:ForwardDiff.Dual{Tag, V, P}, N},
@@ -48,3 +48,82 @@ for op in [:conv, :depthwiseconv, :∇conv_data, :∇conv_filter]
         return ForwardDiff.Dual{Tag, eltype(y), P}.(y, partials)
     end
 end
+
+for op in (:logsoftmax, :softmax)
+    dual_op = Symbol(op, :_dual)
+    @eval function NNlib.$(op)(
+            x::AbstractArray{<:ForwardDiff.Dual{Tag, T, P}}; dims=1) where {Tag, T, P}
+        return Impl.$(dual_op)(x; dims)
+    end
+end
+
+function softmax_dual(
+        x::AbstractArray{<:ForwardDiff.Dual{Tag, T, P}}; dims=1) where {Tag, T, P}
+    value_fn(x) = ForwardDiff.value(Tag, x)
+    partial_fn(x, i) = ForwardDiff.partials(Tag, x, i)
+
+    x_data = value_fn.(x)
+
+    y = NNlib.softmax(x_data; dims)
+    dysᵢ = ntuple(P) do i
+        v = partial_fn.(x, i)
+        return y .* (v .- sum(y .* v; dims))
+    end
+
+    partials = ForwardDiff.Partials.(tuple.(dysᵢ...))
+    return ForwardDiff.Dual{Tag, eltype(y), P}.(y, partials)
+end
+
+function logsoftmax_dual(
+        x::AbstractArray{<:ForwardDiff.Dual{Tag, T, P}}; dims=1) where {Tag, T, P}
+    value_fn(x) = ForwardDiff.value(Tag, x)
+    partial_fn(x, i) = ForwardDiff.partials(Tag, x, i)
+
+    x_data = value_fn.(x)
+
+    y = NNlib.softmax(x_data; dims)
+    dysᵢ = ntuple(P) do i
+        v = partial_fn.(x, i)
+        return v .- sum(y .* v; dims)
+    end
+
+    partials = ForwardDiff.Partials.(tuple.(dysᵢ...))
+    return ForwardDiff.Dual{Tag, eltype(y), P}.(y, partials)
+end
+
+@eval function NNlib.meanpool(
+        x::AbstractArray{<:ForwardDiff.Dual{Tag, T, P}}, pdims::NNlib.PoolDims;
+        kwargs...) where {Tag, T, P}
+    value_fn(x) = ForwardDiff.value(Tag, x)
+    partial_fn(x, i) = ForwardDiff.partials(Tag, x, i)
+
+    y = NNlib.meanpool(value_fn.(x), pdims; kwargs...)
+    dysᵢ = ntuple(P) do i
+        return NNlib.meanpool(partial_fn.(x, i), pdims; kwargs...)
+    end
+
+    partials = ForwardDiff.Partials.(tuple.(dysᵢ...))
+    return ForwardDiff.Dual{Tag, eltype(y), P}.(y, partials)
+end
+
+function NNlib.∇meanpool(
+        dy::AbstractArray{<:ForwardDiff.Dual{Tag, T1, P}},
+        y::AbstractArray{<:ForwardDiff.Dual{Tag, T1, P}},
+        x::AbstractArray{<:ForwardDiff.Dual{Tag, T2, P}},
+        pdims::NNlib.PoolDims; kwargs...) where {Tag, T1, T2, P}
+    value_fn(x) = ForwardDiff.value(Tag, x)
+    partial_fn(x, i) = ForwardDiff.partials(Tag, x, i)
+
+    dy_data, y_data, x_data = value_fn.(dy), value_fn.(y), value_fn.(x)
+
+    dx = NNlib.∇meanpool(dy_data, y_data, x_data, pdims; kwargs...)
+    dysᵢ = ntuple(P) do i
+        ∇y₁ = NNlib.∇meanpool(partial_fn.(dy, i), y_data, x_data, pdims; kwargs...)
+        ∇y₂ = NNlib.∇meanpool(dy_data, partial_fn.(y, i), x_data, pdims; kwargs...)
+        @. ∇y₁ = (∇y₁ + ∇y₂) * partial_fn(x, i)
+        return ∇y₁
+    end
+
+    partials = ForwardDiff.Partials.(tuple.(dysᵢ...))
+    return ForwardDiff.Dual{Tag, eltype(dx), P}.(dx, partials)
+end
diff --git a/lib/LuxLib/test/Project.toml b/lib/LuxLib/test/Project.toml
@@ -39,7 +39,7 @@ BenchmarkTools = "1.5"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.18"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"

diff --git a/lib/LuxLib/test/others/forwarddiff_tests.jl b/lib/LuxLib/test/others/forwarddiff_tests.jl
@@ -26,15 +26,20 @@
 
     function test_jvp_computation(f::F, x, u, ongpu, nested=false) where {F}
         jvp₁ = jvp_forwarddiff(f, x, u)
+
         if !(x isa ComponentArray && ongpu)
             # 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)
+            @testset "JVP ForwardDiff Concrete" begin
+                jvp₂ = jvp_forwarddiff_concrete(f, x, u)
+                @test check_approx(jvp₁, jvp₂; atol=1e-5, rtol=1e-5)
+            end
         end
 
         if !nested
-            jvp₃ = jvp_zygote(f, x, u)
-            @test check_approx(jvp₁, jvp₃; atol=1e-5, rtol=1e-5)
+            @testset "JVP Zygote" begin
+                jvp₃ = jvp_zygote(f, x, u)
+                @test check_approx(jvp₁, jvp₃; atol=1e-5, rtol=1e-5)
+            end
         end
     end
 
@@ -89,6 +94,51 @@
                     true)
             end
         end
+
+        @testset for op in (logsoftmax, softmax)
+            @testset for (input_dim, dim) in zip(
+                (
+                    (2, 3), (2, 3), (2, 3, 4, 5),
+                    (2, 3, 4, 5), (2, 3, 4, 5), (2, 3, 4, 5)
+                ),
+                (1, 2, 1, 2, 3, 4)
+            )
+                x = randn(Float32, input_dim) |> aType
+                u = randn(Float32, input_dim) |> aType
+
+                test_jvp_computation(x -> op(x; dims=dim), x, u, ongpu)
+                test_jvp_computation(
+                    x -> op(x; dims=dim), ComponentArray(; x), u, ongpu)
+
+                test_jvp_computation(
+                    x -> only(Zygote.gradient(x -> sum(op(x; dims=dim)), x)),
+                    x, u, ongpu, true
+                )
+            end
+        end
+
+        @testset for op in (meanpool,)
+            @testset for (input_dim, kernel_size, stride, pad) in (
+                ((8, 3, 2), (4,), (2,), (0,)),
+                ((8, 3, 2), (4,), (3,), (0,)),
+                ((8, 3, 2), (4,), (3,), (1,)),
+                ((8, 8, 3, 2), (4, 4), (2, 2), (0, 0)),
+                ((8, 8, 3, 2), (4, 4), (3, 3), (0, 0)),
+                ((8, 8, 3, 2), (4, 4), (3, 3), (1, 1))
+            )
+                x = randn(Float32, input_dim) |> aType
+                u = randn(Float32, input_dim) |> aType
+
+                test_jvp_computation(
+                    x -> op(x, kernel_size; stride, pad), x, u, ongpu)
+
+                test_jvp_computation(
+                    x -> only(Zygote.gradient(
+                        x -> sum(op(x, kernel_size; stride, pad)), x)),
+                    x, u, ongpu, true
+                )
+            end
+        end
     end
 end
 

diff --git a/test/layers/normalize_tests.jl b/test/layers/normalize_tests.jl
@@ -154,8 +154,6 @@ end
             @jet __f(z)
         end
 
-        broken_backends = VERSION ≥ v"1.11-" ? Any[AutoEnzyme()] : []
-
         @testset "Conv" begin
             c = Conv((3, 3), 3 => 3; init_bias=Lux.ones32)
 
@@ -165,35 +163,31 @@ end
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight, :bias), (2, 2))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight,), (2,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
         end
 
         @testset "Dense" begin
@@ -205,35 +199,31 @@ end
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight, :bias), (2, 2))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight,), (2,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
         end
 
         # See https://github.com/LuxDL/Lux.jl/issues/95