Support for Zygote and ReverseDiff gradients

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: ["1.4", "1.6"]
+        julia-version: ["1.5", "1.6", "~1.7.0-0"]
         os: [ubuntu-latest, macOS-latest]
     steps:
       - uses: actions/checkout@v2

Project.toml

@@ -1,7 +1,7 @@
 name = "Manifolds"
 uuid = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 authors = ["Seth Axen <[email protected]>", "Mateusz Baran <[email protected]>", "Ronny Bergmann <[email protected]>", "Antoine Levitt <[email protected]>"]
-version = "0.6.7"
+version = "0.6.8"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
@@ -42,7 +42,7 @@ SimpleWeightedGraphs = "1"
 SpecialFunctions = "0.8, 0.9, 0.10, 1.0"
 StaticArrays = "1.0"
 StatsBase = "0.32, 0.33"
-julia = "1.4"
+julia = "1.5"
 
 [extras]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
@@ -62,6 +62,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 VisualRegressionTests = "34922c18-7c2a-561c-bac1-01e79b2c4c92"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Colors", "DoubleFloats", "FiniteDiff", "ForwardDiff", "Gtk", "ImageIO", "ImageMagick", "OrdinaryDiffEq", "NLsolve", "Plots", "PyPlot", "Quaternions", "QuartzImageIO", "RecipesBase", "ReverseDiff"]
+test = ["Test", "Colors", "DoubleFloats", "FiniteDiff", "ForwardDiff", "Gtk", "ImageIO", "ImageMagick", "OrdinaryDiffEq", "NLsolve", "Plots", "PyPlot", "Quaternions", "QuartzImageIO", "RecipesBase", "ReverseDiff", "Zygote"]

src/Manifolds.jl

@@ -158,8 +158,8 @@ using RecursiveArrayTools: ArrayPartition
 include("utils.jl")
 
 include("product_representations.jl")
-include("differentiation.jl")
-include("riemannian_diff.jl")
+include("differentiation/differentiation.jl")
+include("differentiation/riemannian_diff.jl")
 
 # Main Meta Manifolds
 include("manifolds/ConnectionManifold.jl")
@@ -285,12 +285,12 @@ end
 function __init__()
     @require FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" begin
         using .FiniteDiff
-        include("finite_diff.jl")
+        include("differentiation/finite_diff.jl")
     end
 
     @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin
         using .ForwardDiff
-        include("forward_diff.jl")
+        include("differentiation/forward_diff.jl")
     end
 
     @require OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" begin
@@ -303,6 +303,11 @@ function __init__()
         include("nlsolve.jl")
     end
 
+    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
+        using .ReverseDiff: ReverseDiff
+        include("differentiation/reverse_diff.jl")
+    end
+
     @require Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" begin
         using .Test: Test
         include("tests/tests_general.jl")
@@ -333,6 +338,12 @@ function __init__()
             include("recipes.jl")
         end
     end
+
+    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
+        using .Zygote: Zygote
+        include("differentiation/zygote.jl")
+    end
+
     return nothing
 end
 

src/differentiation/reverse_diff.jl

@@ -0,0 +1,11 @@
+struct ReverseDiffBackend <: AbstractDiffBackend end
+
+function Manifolds._gradient(f, p, ::ReverseDiffBackend)
+    return ReverseDiff.gradient(f, p)
+end
+
+function Manifolds._gradient!(f, X, p, ::ReverseDiffBackend)
+    return ReverseDiff.gradient!(X, f, p)
+end
+
+push!(Manifolds._diff_backends, ReverseDiffBackend())

src/differentiation/zygote.jl

@@ -0,0 +1,11 @@
+struct ZygoteDiffBackend <: AbstractDiffBackend end
+
+function Manifolds._gradient(f, p, ::ZygoteDiffBackend)
+    return Zygote.gradient(f, p)[1]
+end
+
+function Manifolds._gradient!(f, X, p, ::ZygoteDiffBackend)
+    return copyto!(X, Zygote.gradient(f, p)[1])
+end
+
+push!(Manifolds._diff_backends, ZygoteDiffBackend())

src/nlsolve.jl

@@ -55,6 +55,7 @@ function _inverse_retract_nlsolve(
         retract!(M, F, p, project(M, p, X), retraction; kwargs...)
         project_point && project!(M, q, q)
         F .-= q
+        project_point && project!(M, F, F)
         return F
     end
     isdefined(Manifolds, :NLsolve) ||

test/ambiguities.jl

@@ -4,12 +4,26 @@
         # Interims solution until we follow what was proposed in
         # https://discourse.julialang.org/t/avoid-ambiguities-with-individual-number-element-identity/62465/2
         fmbs = filter(x -> !any(has_type_in_signature.(x, Identity)), mbs)
-        @test length(fmbs) <= 20
+        FMBS_LIMIT = 20
+        @test length(fmbs) <= FMBS_LIMIT
+        if length(fmbs) > FMBS_LIMIT
+            for amb in fmbs
+                println(amb)
+                println()
+            end
+        end
         ms = Test.detect_ambiguities(Manifolds)
         # Interims solution until we follow what was proposed in
         # https://discourse.julialang.org/t/avoid-ambiguities-with-individual-number-element-identity/62465/2
         fms = filter(x -> !any(has_type_in_signature.(x, Identity)), ms)
-        @test length(fms) <= 17
+        FMS_LIMIT = 21
+        if length(fms) > FMS_LIMIT
+            for amb in fms
+                println(amb)
+                println()
+            end
+        end
+        @test length(fms) <= FMS_LIMIT
         # this test takes way too long to perform regularly
         # @test length(our_base_ambiguities()) <= 4
     else

test/differentiation.jl

@@ -17,7 +17,7 @@ using LinearAlgebra: Diagonal, dot
     fd51 = Manifolds.FiniteDifferencesBackend()
     @testset "diff_backend" begin
         @test diff_backend() isa Manifolds.FiniteDifferencesBackend
-        @test length(diff_backends()) == 2
+        @test length(diff_backends()) == 3
         @test diff_backends()[1] isa Manifolds.FiniteDifferencesBackend
 
         @test length(fd51.method.grid) == 5
@@ -33,7 +33,7 @@ using LinearAlgebra: Diagonal, dot
     fwd_diff = Manifolds.ForwardDiffBackend()
     @testset "ForwardDiff" begin
         @test diff_backend() isa Manifolds.FiniteDifferencesBackend
-        @test length(diff_backends()) == 2
+        @test length(diff_backends()) == 3
         @test diff_backends()[1] isa Manifolds.FiniteDifferencesBackend
         @test diff_backends()[2] == fwd_diff
 
@@ -52,8 +52,8 @@ using LinearAlgebra: Diagonal, dot
     finite_diff = Manifolds.FiniteDiffBackend()
     @testset "FiniteDiff" begin
         @test diff_backend() isa Manifolds.FiniteDifferencesBackend
-        @test length(diff_backends()) == 3
-        @test diff_backends()[3] == finite_diff
+        @test length(diff_backends()) == 4
+        @test diff_backends()[4] == finite_diff
 
         @test diff_backend!(finite_diff) == finite_diff
         @test diff_backend() == finite_diff
@@ -65,6 +65,42 @@ using LinearAlgebra: Diagonal, dot
         diff_backend!(fd51)
     end
 
+    using ReverseDiff
+
+    reverse_diff = Manifolds.ReverseDiffBackend()
+    @testset "ReverseDiff" begin
+        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
+        @test length(diff_backends()) == 4
+        @test diff_backends()[3] == reverse_diff
+
+        @test diff_backend!(reverse_diff) == reverse_diff
+        @test diff_backend() == reverse_diff
+        @test diff_backend!(fd51) isa Manifolds.FiniteDifferencesBackend
+        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
+
+        diff_backend!(reverse_diff)
+        @test diff_backend() == reverse_diff
+        diff_backend!(fd51)
+    end
+
+    using Zygote: Zygote
+
+    zygote_diff = Manifolds.ZygoteDiffBackend()
+    @testset "Zygote" begin
+        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
+        @test length(diff_backends()) == 5
+        @test diff_backends()[5] == zygote_diff
+
+        @test diff_backend!(zygote_diff) == zygote_diff
+        @test diff_backend() == zygote_diff
+        @test diff_backend!(fd51) isa Manifolds.FiniteDifferencesBackend
+        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
+
+        diff_backend!(zygote_diff)
+        @test diff_backend() == zygote_diff
+        diff_backend!(fd51)
+    end
+
     @testset "gradient" begin
         diff_backend!(fd51)
         r2 = Euclidean(2)
@@ -74,11 +110,11 @@ using LinearAlgebra: Diagonal, dot
         f2(x) = 3 * x[1] * x[2] + x[2]^3
 
         @testset "Inference" begin
-            v = [-1.0, -1.0]
+            X = [-1.0, -1.0]
             @test (@inferred _derivative(c1, 0.0, Manifolds.ForwardDiffBackend())) ≈
                   [1.0, 0.0]
-            @test (@inferred _derivative!(c1, v, 0.0, Manifolds.ForwardDiffBackend())) === v
-            @test v ≈ [1.0, 0.0]
+            @test (@inferred _derivative!(c1, X, 0.0, Manifolds.ForwardDiffBackend())) === X
+            @test X ≈ [1.0, 0.0]
 
             @test (@inferred _derivative(c1, 0.0, finite_diff)) ≈ [1.0, 0.0]
             @test (@inferred _gradient(f1, [1.0, -1.0], finite_diff)) ≈ [1.0, -2.0]
@@ -87,12 +123,16 @@ using LinearAlgebra: Diagonal, dot
         @testset for backend in [fd51, fwd_diff, finite_diff]
             diff_backend!(backend)
             @test _derivative(c1, 0.0) ≈ [1.0, 0.0]
-            v = [-1.0, -1.0]
-            @test _derivative!(c1, v, 0.0) === v
-            @test isapprox(v, [1.0, 0.0])
+            X = [-1.0, -1.0]
+            @test _derivative!(c1, X, 0.0) === X
+            @test isapprox(X, [1.0, 0.0])
+        end
+        @testset for backend in [fd51, fwd_diff, finite_diff, reverse_diff, zygote_diff]
+            diff_backend!(backend)
+            X = [-1.0, -1.0]
             @test _gradient(f1, [1.0, -1.0]) ≈ [1.0, -2.0]
-            @test _gradient!(f1, v, [1.0, -1.0]) === v
-            @test v ≈ [1.0, -2.0]
+            @test _gradient!(f1, X, [1.0, -1.0]) === X
+            @test X ≈ [1.0, -2.0]
         end
         diff_backend!(Manifolds.NoneDiffBackend())
         @testset for backend in [fd51, Manifolds.ForwardDiffBackend()]

test/groups/special_euclidean.jl

@@ -78,6 +78,7 @@ Random.seed!(10)
                     X_pts;
                     test_diff=true,
                     diff_convs=[(), (LeftAction(),), (RightAction(),)],
+                    atol=1e-9,
                 )
             end
         end

test/manifolds/rotations.jl

@@ -116,7 +116,7 @@ include("../utils.jl")
                 point_distributions=[ptd],
                 tvector_distributions=[tvd],
                 basis_types_to_from=basis_types,
-                exp_log_atol_multiplier=20,
+                exp_log_atol_multiplier=80,
                 retraction_atol_multiplier=12,
                 test_inplace=true,
             )