Merge pull request #464 from dfdx/dfdx/literal-pow-2

Add rules for Base.literal_pow

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.8.19"
+version = "0.8.20"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/base.jl

@@ -161,3 +161,15 @@ end
 @scalar_rule round(x) zero(x)
 @scalar_rule floor(x) zero(x)
 @scalar_rule ceil(x) zero(x)
+
+# note: rules for ^ are defined in the fastmath_able.jl
+function frule((_, _, Δx, _), ::typeof(Base.literal_pow), ::typeof(^), x::Real, pv::Val{p}) where p
+    y = Base.literal_pow(^, x, pv)
+    return y, (p * y / x * Δx)
+end
+
+function rrule(::typeof(Base.literal_pow), ::typeof(^), x::Real, pv::Val{p}) where p
+    y = Base.literal_pow(^, x, pv)
+    literal_pow_pullback(dy) = NoTangent(), NoTangent(), (p * y / x * dy), NoTangent()
+    return y, literal_pow_pullback
+end

test/rulesets/Base/base.jl

@@ -176,4 +176,10 @@
         @test frule((NoTangent(), NoTangent(), NoTangent()), Base.depwarn, "message", :f) !== nothing
         @test rrule(Base.depwarn, "message", :f) !== nothing
     end
+
+    @testset "literal_pow" begin
+        # for real x and n, x must be >0
+        test_frule(Base.literal_pow, ^, 3.5, Val(3))
+        test_rrule(Base.literal_pow, ^, 3.5, Val(3))
+    end
 end