Merge pull request #111 from invenia/mz/densification

Fix _densification_pullback bug

Author: mzgubic

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.34"
+version = "0.1.35"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/chainrules.jl

@@ -25,7 +25,7 @@ function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where
 end
 
 # densification
-function _densification_pullback(Ȳ::Matrix, T, nrows, ncols)
+function _densification_pullback(Ȳ::AbstractMatrix, T, nrows, ncols)
     row_idxs = cumsum(nrows) .- nrows .+ 1
     col_idxs = cumsum(ncols) .- ncols .+ 1
 

test/chainrules.jl

@@ -14,8 +14,9 @@
     end
 
     @testset "Matrix" begin
-        D = BlockDiagonal([randn(1, 2), randn(2, 2)])
+        D = BlockDiagonal([randn(2, 2), randn(2, 2)])
         test_rrule(Matrix, D)
+        test_rrule(Matrix, D, output_tangent=UpperTriangular(rand(4, 4)))
     end
 
     @testset "BlockDiagonal * Vector" begin