Add Enzyme autodiff tests for forward mode in autodiff.jl

dominic-chang · Oct 29, 2024 · c1ad3e1 · c1ad3e1
1 parent 51d50e9
commit c1ad3e1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/test/autodiff.jl b/test/autodiff.jl
@@ -36,6 +36,7 @@ using SpecialFunctions
             @test Zygote.gradient(ϕ -> _F(ϕ, m), ϕ)[1] ≈ 1 / √(1 - m*sin(ϕ)^2) atol=1e-5
             @test ForwardDiff.derivative(ϕ -> _F(ϕ, m), ϕ) ≈ 1 / √(1 - m*sin(ϕ)^2) atol=1e-5
             @test Enzyme.autodiff(Reverse, ϕ -> _F(ϕ, m), Active, Active(ϕ))[1][1] ≈ 1 / √(1 - m*sin(ϕ)^2) atol=1e-5
+            @test Enzyme.autodiff(Forward, ϕ -> _F(ϕ, m), Duplicated, Duplicated(ϕ, 1.0))[1][1] ≈ 1 / √(1 - m*sin(ϕ)^2) atol=1e-5
 
             # 4. ∂m(F(ϕ, m)) == E(ϕ, m) / (2 * m * (1 - m)) - F(ϕ, m) / 2m - sin(2ϕ) / (4 * (1-m) * √(1 - m * sin(ϕ)^2))
             @test Zygote.gradient(m -> _F(ϕ, m), m)[1] ≈
@@ -56,6 +57,7 @@ using SpecialFunctions
             @test Zygote.gradient(ϕ -> _E(ϕ, m), ϕ)[1] ≈ √(1 - m * sin(ϕ)^2) atol=1e-5
             @test ForwardDiff.derivative(ϕ -> _E(ϕ, m), ϕ) ≈ √(1 - m * sin(ϕ)^2) atol=1e-5
             @test Enzyme.autodiff(Reverse, ϕ -> _E(ϕ, m), Active, Active(ϕ))[1][1] ≈ √(1 - m * sin(ϕ)^2) atol=1e-5
+            @test Enzyme.autodiff(Forward, ϕ -> _E(ϕ, m), Duplicated, Duplicated(ϕ, 1.0))[1][1] ≈ √(1 - m*sin(ϕ)^2) atol=1e-5
 
             # 6. ∂m(E(ϕ, m)) == (E(ϕ, m) - F(ϕ, m)) / 2m
             @test Zygote.gradient(m -> _E(ϕ, m), m)[1] ≈ (alg.E(ϕ, m) - alg.F(ϕ, m)) / 2m