Add ChainRules support (#202)

* Add ChainRules support

* Add chain rule for differentiate

* Add chain rule for -

* Add chain rule for *

Author: blegat

Project.toml

@@ -5,6 +5,7 @@ repo = "https://github.com/JuliaAlgebra/MultivariatePolynomials.jl"
 version = "0.4.4"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"

src/MultivariatePolynomials.jl

@@ -88,4 +88,6 @@ include("division.jl")
 include("gcd.jl")
 include("det.jl")
 
+include("chain_rules.jl")
+
 end # module

src/chain_rules.jl

@@ -0,0 +1,27 @@
+import ChainRulesCore
+
+ChainRulesCore.@scalar_rule +(x::APL, y::APL) (true, true)
+ChainRulesCore.@scalar_rule -(x::APL, y::APL) (true, -1)
+
+function ChainRulesCore.frule((_, Δp, Δq), ::typeof(*), p::APL, q::APL)
+    return p * q, MA.add_mul!!(p * Δq, q, Δp)
+end
+function ChainRulesCore.rrule(::typeof(*), p::APL, q::APL)
+    function times_pullback2(ΔΩ̇)
+        #ΔΩ = ChainRulesCore.unthunk(Ω̇)
+        #return (ChainRulesCore.NoTangent(), ChainRulesCore.ProjectTo(p)(ΔΩ * q'), ChainRulesCore.ProjectTo(q)(p' * ΔΩ))
+        return (ChainRulesCore.NoTangent(), ΔΩ̇ * q', p' * ΔΩ̇)
+    end
+    return p * q, times_pullback2
+end
+
+function ChainRulesCore.frule((_, Δp, _), ::typeof(differentiate), p, x)
+    return differentiate(p, x), differentiate(Δp, x)
+end
+function pullback(Δdpdx, x)
+    return ChainRulesCore.NoTangent(), x * differentiate(x * Δdpdx, x), ChainRulesCore.NoTangent()
+end
+function ChainRulesCore.rrule(::typeof(differentiate), p, x)
+    dpdx = differentiate(p, x)
+    return dpdx, Base.Fix2(pullback, x)
+end

test/chain_rules.jl

@@ -0,0 +1,57 @@
+using LinearAlgebra, Test
+using ChainRulesCore
+
+function test_chain_rule(dot, op, args, Δin, Δout)
+    output = op(args...)
+    foutput, fΔout = ChainRulesCore.frule((NoTangent(), Δin...), op, args...)
+    @test output == foutput
+    routput, pullback = ChainRulesCore.rrule(op, args...)
+    @test output == routput
+    rΔin = pullback(Δout)
+    @test rΔin[1] == NoTangent()
+    @test dot(Δin, rΔin[2:end]) ≈ dot(fΔout, Δout)
+end
+
+@testset "ChainRulesCore" begin
+    Mod.@polyvar x y
+    p = 1.1x + y
+    q = (-0.1 + im) * x - y
+
+    output, pullback = ChainRulesCore.rrule(+, p, q)
+    @test output == (1.0 + im)x
+    @test pullback(2) == (NoTangent(), 2, 2)
+    @test pullback(x + 3) == (NoTangent(), x + 3, x + 3)
+
+    output, pullback = ChainRulesCore.rrule(-, p, q)
+    @test output ≈ (1.2 - im) * x + 2y
+    @test pullback(2) == (NoTangent(), 2, -2)
+    @test pullback(x + 3) == (NoTangent(), x + 3, -x - 3)
+
+    output, pullback = ChainRulesCore.rrule(differentiate, p, x)
+    @test output == 1.1
+    @test pullback(q) == (NoTangent(), (-0.2 + 2im) * x^2 - x*y, NoTangent())
+    @test pullback(1x) == (NoTangent(), 2x^2, NoTangent())
+
+    test_chain_rule(dot, +, (p, q), (q, p), p)
+    test_chain_rule(dot, +, (p, q), (p, q), q)
+
+    test_chain_rule(dot, -, (p, q), (q, p), p)
+    test_chain_rule(dot, -, (p, q), (p, q), q)
+
+    test_chain_rule(dot, *, (p, q), (q, p), p * q)
+    test_chain_rule(dot, *, (p, q), (p, q), q * q)
+    test_chain_rule(dot, *, (q, p), (p, q), q * q)
+    test_chain_rule(dot, *, (p, q), (q, p), q * q)
+
+    function _dot(p, q)
+        monos = monomials(p + q)
+        return dot(coefficient.(p, monos), coefficient.(q, monos))
+    end
+    function _dot(px::Tuple{<:AbstractPolynomial,NoTangent}, qx::Tuple{<:AbstractPolynomial,NoTangent})
+        return _dot(px[1], qx[1])
+    end
+    test_chain_rule(_dot, differentiate, (p, x), (q, NoTangent()), p)
+    test_chain_rule(_dot, differentiate, (p, x), (q, NoTangent()), differentiate(p, x))
+    test_chain_rule(_dot, differentiate, (p, x), (q, NoTangent()), differentiate(q, x))
+    test_chain_rule(_dot, differentiate, (p, x), (p * q, NoTangent()), p)
+end

test/commutativetests.jl

@@ -20,6 +20,7 @@ include("comparison.jl")
 include("substitution.jl")
 include("differentiation.jl")
 include("division.jl")
+include("chain_rules.jl")
 
 include("show.jl")