Merge pull request #122 from invenia/ed/fix-diagonal

Fix Diagonal sensitivity
invenia · Jan 8, 2019 · 30d5ccc · 30d5ccc
2 parents b510378 + 004951f
commit 30d5ccc
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Nabla"
 uuid = "49c96f43-aa6d-5a04-a506-44c7070ebe78"
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"

diff --git a/src/sensitivities/linalg/diagonal.jl b/src/sensitivities/linalg/diagonal.jl
@@ -64,21 +64,21 @@ function ∇(
     ::Type{Arg{1}},
     p,
     Y::∇ScalarDiag,
-    Ȳ::∇ScalarDiag,
+    Ȳ::∇AbstractMatrix,
     x::∇AbstractVector,
 )
-    return copyto!(similar(x), Ȳ.diag)
+    return copyto!(similar(x), diag(Ȳ))
 end
 function ∇(
     x̄::∇AbstractVector,
     ::Type{Diagonal},
     ::Type{Arg{1}},
     p,
     Y::∇ScalarDiag,
-    Ȳ::∇ScalarDiag,
+    Ȳ::∇AbstractMatrix,
     x::∇AbstractVector,
 )
-    return broadcast!(+, x̄, x̄, Ȳ.diag)
+    return broadcast!(+, x̄, x̄, diag(Ȳ))
 end
 
 @explicit_intercepts Diagonal Tuple{∇AbstractMatrix}

diff --git a/test/sensitivities/linalg/diagonal.jl b/test/sensitivities/linalg/diagonal.jl
@@ -56,4 +56,13 @@
             @test check_errs(logdet, 10.0, A, VA)
         end
     end
+
+    # test that both diagonal implementations are correct and the same
+    let
+        f1(x) = sum(Diagonal(x))
+        f2(x) = sum(diagm(0 => x))
+
+        @test ∇(f1)(ones(4))[1] == ones(4)
+        @test ∇(f2)(ones(4))[1] == ones(4)
+    end
 end