tidy

Author: mcabbott

src/rulesets/Base/fastmath_able.jl

@@ -167,38 +167,44 @@ let
         # literal_pow is in base.jl
         function frule((_, Δx, Δp), ::typeof(^), x::Number, p::Number)
             y = x ^ p
-            dx = _pow_grad_x(x, p, float(y))
+            _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)
+            _dp = Δp isa AbstractZero ? Δp : _pow_grad_p(x, p, float(y))
+            return y, muladd(_dp, Δp, _dx * Δx)
         end
 
         function rrule(::typeof(^), x::Number, p::Number)
             y = x^p
-            project_x, project_p = ProjectTo(x), ProjectTo(p)
+            project_x = ProjectTo(x)
+            project_p = ProjectTo(p)
             @inline function power_pullback(dy)
-                dx = project_x(conj(_pow_grad_x(x,p,float(y))) * dy)
-                dp = @thunk project_p(conj(_pow_grad_p(x,p,float(y))) * dy)
-                return (NoTangent(), dx, dp)
+                _dx = _pow_grad_x(x, p, float(y))
+                _dy = _pow_grad_p(x, p, float(y))
+                return (
+                    NoTangent(), 
+                    project_x(conj(_dx) * dy),
+                    @thunk project_p(conj(_dy) * dy)
+                    )
             end
             return y, power_pullback
         end
 
+        ## `rem`
         @scalar_rule(
             rem(x, y),
             @setup((u, nan) = promote(x / y, NaN16), isint = isinteger(x / y)),
             (ifelse(isint, nan, one(u)), ifelse(isint, nan, -trunc(u))),
         )
+        ## `min`, `max`
         @scalar_rule max(x, y) @setup(gt = x > y) (gt, !gt)
         @scalar_rule min(x, y) @setup(gt = x > y) (!gt, gt)
 
         # Unary functions
         @scalar_rule +x true
         @scalar_rule -x -1
 
-        # `sign`
-
+        ## `sign`
         function frule((_, Δx), ::typeof(sign), x)
             n = ifelse(iszero(x), one(real(x)), abs(x))
             Ω = x isa Real ? sign(x) : x / n
@@ -263,11 +269,6 @@ end
 # Thes functions need to be defined outside the eval() block.
 # The special cases they aim to hit are in POWERGRADS in tests.
 _pow_grad_x(x, p, y) = (p * y / x)
-# function _pow_grad_x(x::Real, p::Real, y)
-#     return ifelse(!iszero(x) | (p<0), (p * y / x),
-#              ifelse(isone(p), one(y),
-#                ifelse((0<p) | (p<1), oftype(y, Inf), zero(y) )))
-# end
 function _pow_grad_x(x::Real, p::Real, y)
     return if !iszero(x) || p < 0
         p * y / x
@@ -281,10 +282,6 @@ function _pow_grad_x(x::Real, p::Real, y)
 end
 
 _pow_grad_p(x, p, y) = y * log(complex(x))
-# function _pow_grad_p(x::Real, p::Real, y)
-#     return ifelse(!iszero(x), y * real(log(complex(x))),
-#              ifelse(p>0, zero(y), oftype(y, NaN) ))
-# end
 function _pow_grad_p(x::Real, p::Real, y)
     return if !iszero(x)
         y * real(log(complex(x)))

test/rulesets/Base/fastmath_able.jl

@@ -180,7 +180,7 @@ const FASTABLE_AST = quote
         # At x=0, gradients for x seem clear, 
         # for p less certain what's best.
           (0.0, 2)   => (0.0, 0.0),
-          (-0.0, 2)  => (-0.0, 0.0),
+          (-0.0, 2)  => (0.0, 0.0),  # probably (-0.0, 0.0) would be ideal
           (0.0, 1)   => (1.0, 0.0),
           (-0.0, 1)  => (1.0, 0.0),
           (0.0, 0)   => (0.0, NaN),
@@ -210,7 +210,7 @@ const FASTABLE_AST = quote
 
             # Forward
             y_fwd = frule((1,1,1), ^, x, p)[1]
-            @test y === y_fwd # || println("^ forward value for $x^$p: got $y_fwd, expected $y")
+            @test isequal(y, y_fwd)
 
             # ∂x_fwd = frule((0,1,0), ^, x, p)[1]
             # ∂p_fwd = frule((0,0,1), ^, x, p)[2]
@@ -219,15 +219,11 @@ const FASTABLE_AST = quote
 
             # Reverse
             y_rev = rrule(^, x, p)[1]
-            @test y === y_rev # || println("^ reverse value for $x^$p: got $y_rev, expected $y")
+            @test isequal(y, y_rev)
 
             ∂x_rev, ∂p_rev = unthunk.(rrule(^, x, p)[2](1))[2:3]
-            if ∂x === -0.0 # happens at at x === -0.0 && p === 2, ignore the sign
-                @test 0.0 == ∂x_rev
-            else
-                @test isequal(∂x, ∂x_rev) # || println("^ reverse `x` gradient for $x^$p: got $∂x_rev, expected $∂x")
-            end
-            @test isequal(∂p, ∂p_rev) # || println("^ reverse `p` gradient for $x^$p: got $∂p_rev, expected $∂p")
+            @test isequal(∂x, ∂x_rev)
+            @test isequal(∂p, ∂p_rev)
         end
     end