Add new-style diagm to tracker

2018-10-26 13:39:49 -05:00 · 2018-10-26 13:39:49 -05:00 · 9f9803eec6
commit 9f9803eec6
parent 44ccdb7ca9
1 changed files with 3 additions and 0 deletions
--- a/src/tracker/array.jl
+++ b/src/tracker/array.jl
@ -309,6 +309,9 @@ end

 # BLAS

+LinearAlgebra.diagm(x::Pair{<:Integer, <:TrackedVector}) = track(diagm, x)
+@grad diagm(x::Pair) = diagm(x[1] => data(x[2])), Δ -> (diag(Δ, x[1]),)
+
 LinearAlgebra.diagm(x::TrackedVector) = track(diagm, x)
@grad diagm(x) = diagm(data(x)), Δ -> (diag(Δ),)