mostly delete literal_pow

Author: mcabbott

src/rulesets/Base/base.jl

@@ -180,26 +180,23 @@ end
 @scalar_rule ceil(x) zero(x)
 
 # `literal_pow`
+# This is mostly handled by AD; it's a micro-optimisation to provide a gradient for x*x*x
 # Note that rules for `^` are defined in the fastmath_able.jl
 
-function frule((_, _, Δx, _), ::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{p}) where p
-    y = Base.literal_pow(^, x, Val(p))
-    yox = Base.literal_pow(^, x, Val(p-1))
-    return y, p * yox * Δx
+function frule((_, _, Δx, _), ::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{2})
+    return x * x, 2 * x * Δx
+end
+function frule((_, _, Δx, _), ::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{3})
+    x2 = x * x
+    return x2 * x, 3 * x2 * Δx
 end
-frule((_, _, Δx, _), ::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{0}) = x^0, zero(Δx)
 
-function rrule(::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{p}) where p
-    y = Base.literal_pow(^, x, Val(p))
-    @inline function literal_pow_pullback(dy)
-        # Calling literal_pow a 2nd time is the easy way to get all the edge cases right.
-        # It should be cheap up to p=4, which is the main use of literal powers, right?
-        yox = Base.literal_pow(^, x, Val(p-1))
-        return (NoTangent(), NoTangent(), ProjectTo(x)(p * yox * dy), NoTangent())
-    end
-    return y, literal_pow_pullback
+function rrule(::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{2})
+    @inline pow2_pullback(dy) = (NoTangent(), NoTangent(), ProjectTo(x)(2 * x * dy), NoTangent())
+    return x * x, pow2_pullback
 end
-function rrule(::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{0})
-    literal_pow_zero_pullback(dy) = (NoTangent(), NoTangent(), ProjectTo(x)(zero(dy)), NoTangent())
-    return x^0, literal_pow_zero_pullback
+function rrule(::typeof(Base.literal_pow), ::typeof(^), x::Real, ::Val{3})
+    x2 = x * x
+    @inline pow3_pullback(dy) = (NoTangent(), NoTangent(), ProjectTo(x)(3 * x2 * dy), NoTangent())
+    return x2 * x, pow3_pullback
 end

test/rulesets/Base/base.jl

@@ -188,46 +188,10 @@
         @test rrule(Base.depwarn, "message", :f) !== nothing
     end
 
-    @testset "literal_pow" begin
-        test_frule(Base.literal_pow, ^, 3.5, Val(3))
-        test_rrule(Base.literal_pow, ^, 3.5, Val(3))
-
-        @testset "regular: $x^$p" for x in [-1.5, 0.0, 3.5], p in [-3, -1, 0, 1, 3]
-            x == 0 && p < 0 && continue
-            test_frule(Base.literal_pow, ^, x, Val(p))
-            test_rrule(Base.literal_pow, ^, x, Val(p))
-        end
-
-        @testset "singularities: 0^0, 0^-1, 0^-2" begin
-            # Trivial one: 0^0 == 1 in Julia
-            @test frule((1,1,1,1), Base.literal_pow, ^, 0.0, Val(0)) == ((0.0)^0, 0)
-            @test rrule(Base.literal_pow, ^, 0.0, Val(0))[2](1.0)[3] == 0.0
-            
-            # Odd power, 1/x
-            @test frule((1,1,1,1), Base.literal_pow, ^, 0.0, Val(-1)) == ((0.0)^-1, -Inf)
-            @test rrule(Base.literal_pow, ^, 0.0, Val(-1))[1] == (0.0)^-1 == Inf
-            @test rrule(Base.literal_pow, ^, 0.0, Val(-1))[2](1.0)[3] == -Inf
-
-            @test frule((1,1,1,1), Base.literal_pow, ^, -0.0, Val(-1)) == ((-0.0)^-1, -Inf)
-            @test rrule(Base.literal_pow, ^, -0.0, Val(-1))[1] == (-0.0)^-1 == -Inf
-            @test rrule(Base.literal_pow, ^, -0.0, Val(-1))[2](1.0)[3] == -Inf
-
-            # Even power, 1/x^2
-            @test frule((1,1,1,1), Base.literal_pow, ^, 0.0, Val(-2)) == ((0.0)^-2, -Inf)
-            @test rrule(Base.literal_pow, ^, 0.0, Val(-2))[1] == (0.0)^-2 == Inf
-            @test rrule(Base.literal_pow, ^, 0.0, Val(-2))[2](1.0)[3] == -Inf
-
-            @test frule((1,1,1,1), Base.literal_pow, ^, -0.0, Val(-2)) == ((-0.0)^-2, +Inf)
-            @test rrule(Base.literal_pow, ^, -0.0, Val(-2))[1] == (-0.0)^-2 == Inf
-            @test rrule(Base.literal_pow, ^, -0.0, Val(-2))[2](1.0)[3] == +Inf
-
-            # Not singluar, but ^ messed these up: x^1 and x^2
-            @test frule((1,1,1,1), Base.literal_pow, ^, 0.0, Val(2)) == (0.0, 0)
-            @test rrule(Base.literal_pow, ^, 0.0, Val(2))[2](1.0)[3] == 0.0
-
-            @test frule((1,1,1,1), Base.literal_pow, ^, 0.0, Val(1)) == (0.0, 1)
-            @test rrule(Base.literal_pow, ^, 0.0, Val(1))[2](1.0)[3] == 1.0
-        end
+    @testset "literal_pow: $x^$p" for x in [-1.5, 0.0, 3.5], p in [2, 3]
+        x == 0 && p < 0 && continue
+        test_frule(Base.literal_pow, ^, x, Val(p))
+        test_rrule(Base.literal_pow, ^, x, Val(p))
     end
 
     @testset "Float conversions" begin

test/rulesets/Base/fastmath_able.jl

@@ -172,13 +172,7 @@ const FASTABLE_AST = quote
             test_rrule(^, rand(T), rand(T))
         end
 
-        # @testset "^(x::$T, $p::Int)" for T in (Float64, ComplexF64), p in -2:2
-        #     test_frule(^, randn(T) + 3, p ⊢ NoTangent())  # this doesn't just skip p's tangent
-        #     test_rrule(^, randn(T) + 3, p ⊢ NoTangent())
-        # end
-
         # Tests for power functions, at values near to zero.
-
         POWERGRADS = [ # (x,p) => (dx,dp)
         # Some regular points, as sanity checks:
           (1.0, 2)   => (2.0, 0.0),
@@ -235,24 +229,6 @@ const FASTABLE_AST = quote
             end
             @test isequal(∂p, ∂p_rev) # || println("^ reverse `p` gradient for $x^$p: got $∂p_rev, expected $∂p")
         end
-
-        @testset "literal_pow $x ^ $p" for ((x,p), (∂x, ∂p)) in POWERGRADS
-            y = Base.literal_pow(^, x, Val(p))
-
-            # Forward
-            y_fwd = frule((1,1,1,1), Base.literal_pow, ^, x, Val(p))[1]
-            @test y === y_fwd # || println("literal_pow forward value for $x^$p: got $y_fwd, expected $y")
-
-            ∂x_fwd = frule((0,0,1,0), Base.literal_pow, ^, x, Val(p))[1]
-            # isequal(∂x, ∂x_fwd) || println("literal_pow forward `x` gradient for $x^$p: got $∂x_fwd, expected $∂x, maybe, y=$y")
-
-            # Reverse
-            y_rev = rrule(Base.literal_pow, ^, x, Val(p))[1]
-            @test y === y_rev # || println("literal_pow reverse value for $x^$p: got $y_rev, expected $y")
-
-            ∂x_rev = unthunk(rrule(Base.literal_pow, ^, x, Val(p))[2](1))[3]
-            @test isequal(∂x, ∂x_rev) # || println("literal_pow `x` gradient for $x^$p: got $∂x_rev, expected $∂x")
-        end
     end
 
     @testset "sign" begin