fix frule

Author: mcabbott

src/rulesets/Base/fastmath_able.jl

@@ -168,10 +168,14 @@ let
         function frule((_, Δx, Δp), ::typeof(^), x::Number, p::Number)
             y = x ^ p
             _dx = _pow_grad_x(x, p, float(y))
-            # When x < 0 && isinteger(p), could decide p is non-differentiable, isolated
-            # points, but chose to match what the rrule with ProjectTo gives, real(log(...)):
-            _dp = Δp isa AbstractZero ? Δp : _pow_grad_p(x, p, float(y))
-            return y, muladd(_dp, Δp, _dx * Δx)
+            if iszero(Δp)
+                # Treat this as a strong zero, to avoid NaN, and save the cost of log
+                return y, _dx * Δx
+            else
+                # This may do real(log(complex(...))) which matches ProjectTo in rrule
+                _dp = _pow_grad_p(x, p, float(y))
+                return y, muladd(_dp, Δp, _dx * Δx)
+            end
         end
 
         function rrule(::typeof(^), x::Number, p::Number)

test/rulesets/Base/fastmath_able.jl

@@ -212,10 +212,17 @@ const FASTABLE_AST = quote
             y_fwd = frule((1,1,1), ^, x, p)[1]
             @test isequal(y, y_fwd)
 
-            # ∂x_fwd = frule((0,1,0), ^, x, p)[1]
-            # ∂p_fwd = frule((0,0,1), ^, x, p)[2]
-            # isequal(∂x, ∂x_fwd) || println("^ forward `x` gradient for $y = $x^$p: got $∂x_fwd, expected $∂x, maybe!")
-            # isequal(∂p, ∂p_fwd) || println("^ forward `p` gradient for $x^$p: got $∂p_fwd, expected $∂p, maybe")
+            ∂x_fwd = frule((0,1,0), ^, x, p)[2]
+            ∂p_fwd = frule((0,0,1), ^, x, p)[2]
+            @test isequal(∂x, ∂x_fwd)
+            if x===0.0 && p===0.5
+                @test_broken isequal(∂p, ∂p_fwd)
+            else
+                @test isequal(∂p, ∂p_fwd)
+            end
+
+            ∂x_fwd = frule((0,1,ZeroTangent()), ^, x, p)[2] # easier, strong zero
+            @test isequal(∂x, ∂x_fwd)
 
             # Reverse
             y_rev = rrule(^, x, p)[1]