test wirtinger derivatives with fdm too

Author: simeonschaub

test/rulesets/Base/base.jl

@@ -87,11 +87,13 @@
 
             x isa Real && test_scalar(cbrt, x)
             if (x isa Real && x >= 0) || x isa Complex
-                test_scalar(sqrt, x)
-                test_scalar(log, x)
-                test_scalar(log2, x)
-                test_scalar(log10, x)
-                test_scalar(log1p, x)
+                # this check is needed because the imaginary part jumps from pi to -pi
+                nr = imag(x) != 0 && real(x) < 0
+                test_scalar(sqrt, x, test_wirtinger=nr)
+                test_scalar(log, x, test_wirtinger=nr)
+                test_scalar(log2, x, test_wirtinger=nr)
+                test_scalar(log10, x, test_wirtinger=nr)
+                test_scalar(log1p, x, test_wirtinger=nr)
             end
         end
     end

test/test_util.jl

@@ -16,18 +16,33 @@ at input point `x` to confirm that there are correct ChainRules provided.
 
 All keyword arguments except for `fdm` are passed to `isapprox`.
 """
-function test_scalar(f, x; rtol=1e-9, atol=1e-9, fdm=_fdm, kwargs...)
+function test_scalar(f, x; rtol=1e-9, atol=1e-9, fdm=_fdm, test_wirtinger=x isa Complex, kwargs...)
     @testset "$f at $x, $(nameof(rule))" for rule in (rrule, frule)
         res = rule(f, x)
         @test res !== nothing  # Check the rule was defined
         fx, ∂x = res
         @test fx == f(x)  # Check we still get the normal value, right
 
         # Check that we get the derivative right:
-        @test isapprox(
-            ∂x(1), fdm(f, x);
-            rtol=rtol, atol=atol, kwargs...
-        )
+        if !test_wirtinger
+            @test isapprox(
+                ∂x(1), fdm(f, x);
+                rtol=rtol, atol=atol, kwargs...
+            )
+        else
+            # For complex arguments, also check if the wirtinger derivative is correct
+            ∂Re, ∂Im = fdm.((ϵ -> f(x + ϵ), ϵ -> f(x + im*ϵ)), 0)
+            ∂ = .5(∂Re - im*∂Im)
+            ∂̅ = .5(∂Re + im*∂Im)
+            @test isapprox(
+                wirtinger_primal(∂x(1)), ∂;
+                rtol=rtol, atol=atol, kwargs...
+            )
+            @test isapprox(
+                extern(wirtinger_conjugate(∂x(1))), ∂̅;
+                rtol=rtol, atol=atol, kwargs...
+            )
+        end
     end
 end