Merge 457ed82 into 7b69721

JuliaDiff · Apr 22, 2019 · 24d8eec · 24d8eec
2 parents 7b69721 + 457ed82
commit 24d8eec
Show file tree

Hide file tree

Showing 8 changed files with 34 additions and 2 deletions.
diff --git a/src/ChainRules.jl b/src/ChainRules.jl
@@ -10,7 +10,9 @@ include("differentials.jl")
 include("rules.jl")
 include("rules/base.jl")
 include("rules/broadcast.jl")
-include("rules/linalg.jl")
+include("rules/linalg/dense.jl")
+include("rules/linalg/diagonal.jl")
+include("rules/linalg/symmetric.jl")
 include("rules/blas.jl")
 include("rules/nanmath.jl")
 include("rules/specialfunctions.jl")

diff --git a/src/rules/linalg.jl → src/rules/linalg/dense.jl b/src/rules/linalg.jl → src/rules/linalg/dense.jl
diff --git a/src/rules/linalg/diagonal.jl b/src/rules/linalg/diagonal.jl
@@ -0,0 +1,2 @@
+rrule(::typeof(Diagonal), d::AbstractVector) = Diagonal(d), Rule(diag)
+rrule(::typeof(diag), A::AbstractMatrix) = diag(A), Rule(Diagonal)
diff --git a/src/rules/linalg/symmetric.jl b/src/rules/linalg/symmetric.jl
@@ -0,0 +1,4 @@
+rrule(::typeof(Symmetric), A::AbstractMatrix) = Symmetric(A), Rule(_symmetric_back)
+
+_symmetric_back(ΔΩ) = UpperTriangular(ΔΩ) + LowerTriangular(ΔΩ)' - Diagonal(ΔΩ)
+_symmetric_back(ΔΩ::Union{Diagonal, UpperTriangular}) = ΔΩ
diff --git a/test/rules/linalg.jl → test/rules/linalg/dense.jl b/test/rules/linalg.jl → test/rules/linalg/dense.jl
diff --git a/test/rules/linalg/diagonal.jl b/test/rules/linalg/diagonal.jl
@@ -0,0 +1,14 @@
+@testset "diagonal" begin
+    @testset "Diagonal" begin
+        rng, N = MersenneTwister(123456), 3
+        rrule_test(Diagonal, randn(rng, N, N), (randn(rng, N), randn(rng, N)))
+        rrule_test(Diagonal, Diagonal(randn(rng, N)), (randn(rng, N), randn(rng, N)))
+    end
+    @testset "diag" begin
+        rng, N = MersenneTwister(123456), 7
+        rrule_test(diag, randn(rng, N), (randn(rng, N, N), randn(rng, N, N)))
+        rrule_test(diag, randn(rng, N), (Diagonal(randn(rng, N)), randn(rng, N, N)))
+        rrule_test(diag, randn(rng, N), (randn(rng, N, N), Diagonal(randn(rng, N))))
+        rrule_test(diag, randn(rng, N), (Diagonal(randn(rng, N)), Diagonal(randn(rng, N))))
+    end
+end
diff --git a/test/rules/linalg/symmetric.jl b/test/rules/linalg/symmetric.jl
@@ -0,0 +1,6 @@
+@testset "symmetric" begin
+    @testset "Symmetric" begin
+        rng, N = MersenneTwister(123456), 3
+        rrule_test(Symmetric, randn(rng, N, N), (randn(rng, N, N), randn(rng, N, N)))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -14,8 +14,12 @@ include("test_util.jl")
     include("rules.jl")
     @testset "rules" begin
         include(joinpath("rules", "base.jl"))
+        @testset "linalg" begin
+            include(joinpath("rules", "linalg", "dense.jl"))
+            include(joinpath("rules", "linalg", "diagonal.jl"))
+            include(joinpath("rules", "linalg", "symmetric.jl"))
+        end
         include(joinpath("rules", "broadcast.jl"))
-        include(joinpath("rules", "linalg.jl"))
         include(joinpath("rules", "blas.jl"))
         include(joinpath("rules", "nanmath.jl"))
         include(joinpath("rules", "specialfunctions.jl"))