Add more diag/diagm rrules (#200)

* Add rrule for kth diag

* Add rrules for diagm

* Add version bound for diagm with size

* Increment version number

Project.toml

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

src/rulesets/LinearAlgebra/structured.jl

@@ -25,6 +25,35 @@ function rrule(::typeof(diag), A::AbstractMatrix)
     end
     return diag(A), diag_pullback
 end
+if VERSION ≥ v"1.3"
+    function rrule(::typeof(diag), A::AbstractMatrix, k::Integer)
+        function diag_pullback(ȳ)
+            return (NO_FIELDS, @thunk(diagm(size(A)..., k => ȳ)), DoesNotExist())
+        end
+        return diag(A, k), diag_pullback
+    end
+
+    function rrule(::typeof(diagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)
+        function diagm_pullback(ȳ)
+            return (NO_FIELDS, DoesNotExist(), DoesNotExist(), _diagm_back.(kv, Ref(ȳ))...)
+        end
+        return diagm(m, n, kv...), diagm_pullback
+    end
+end
+function rrule(::typeof(diagm), kv::Pair{<:Integer,<:AbstractVector}...)
+    function diagm_pullback(ȳ)
+        return (NO_FIELDS, _diagm_back.(kv, Ref(ȳ))...)
+    end
+    return diagm(kv...), diagm_pullback
+end
+
+function _diagm_back(p, ȳ)
+    return Thunk() do 
+        k, v = p
+        d = diag(ȳ, k)[1:length(v)] # handle if diagonal was smaller than matrix
+        return Composite{typeof(p)}(second = d)
+    end
+end
 
 function rrule(::typeof(*), D::Diagonal{<:Real}, V::AbstractVector{<:Real})
     function times_pullback(Ȳ)

test/rulesets/LinearAlgebra/structured.jl

@@ -30,6 +30,53 @@
         rrule_test(diag, randn(N), (Diagonal(randn(N)), randn(N, N)))
         rrule_test(diag, randn(N), (randn(N, N), Diagonal(randn(N))))
         rrule_test(diag, randn(N), (Diagonal(randn(N)), Diagonal(randn(N))))
+        VERSION ≥ v"1.3" && @testset "k=$k" for k in (-1, 0, 2)
+            M = N - abs(k)
+            rrule_test(diag, randn(M), (randn(N, N), randn(N, N)), (k, nothing))
+        end
+    end
+    @testset "diagm" begin
+        @testset "without size" begin
+            M, N = 7, 9
+            s = (8, 8)
+            a, ā = randn(M), randn(M)
+            b, b̄ = randn(M), randn(M)
+            c, c̄ = randn(M - 1), randn(M - 1)
+            ȳ = randn(s)
+            ps = (0 => a, 1 => b, 0 => c)
+            y, back = rrule(diagm, ps...)
+            @test y == diagm(ps...)
+            ∂self, ∂pa, ∂pb, ∂pc = back(ȳ)
+            @test ∂self === NO_FIELDS
+            ∂a_fd, ∂b_fd, ∂c_fd = j′vp(_fdm, (a, b, c) -> diagm(0 => a, 1 => b, 0 => c), ȳ, a, b, c)
+            for (p, ∂px, ∂x_fd) in zip(ps, (∂pa, ∂pb, ∂pc), (∂a_fd, ∂b_fd, ∂c_fd))
+                ∂px = unthunk(∂px)
+                @test ∂px isa Composite{typeof(p)}
+                @test ∂px.first isa AbstractZero
+                @test ∂px.second ≈ ∂x_fd
+            end
+        end
+        VERSION ≥ v"1.3" && @testset "with size" begin
+            M, N = 7, 9
+            a, ā = randn(M), randn(M)
+            b, b̄ = randn(M), randn(M)
+            c, c̄ = randn(M - 1), randn(M - 1)
+            ȳ = randn(M, N)
+            ps = (0 => a, 1 => b, 0 => c)
+            y, back = rrule(diagm, M, N, ps...)
+            @test y == diagm(M, N, ps...)
+            ∂self, ∂M, ∂N, ∂pa, ∂pb, ∂pc = back(ȳ)
+            @test ∂self === NO_FIELDS
+            @test ∂M === DoesNotExist()
+            @test ∂N === DoesNotExist()
+            ∂a_fd, ∂b_fd, ∂c_fd = j′vp(_fdm, (a, b, c) -> diagm(M, N, 0 => a, 1 => b, 0 => c), ȳ, a, b, c)
+            for (p, ∂px, ∂x_fd) in zip(ps, (∂pa, ∂pb, ∂pc), (∂a_fd, ∂b_fd, ∂c_fd))
+                ∂px = unthunk(∂px)
+                @test ∂px isa Composite{typeof(p)}
+                @test ∂px.first isa AbstractZero
+                @test ∂px.second ≈ ∂x_fd
+            end
+        end
     end
     @testset "Symmetric" begin
         N = 3