Merge pull request #111 from invenia/mz/densification

Fix _densification_pullback bug
JuliaArrays · Jul 25, 2022 · 6223d3f · 6223d3f · mzgubic · Jul 25, 2022
2 parents 7bac0da + e7f8c20
commit 6223d3f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/Project.toml b/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"

diff --git a/src/chainrules.jl b/src/chainrules.jl
@@ -25,7 +25,7 @@ function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where
 end
 
 # densification
-function _densification_pullback(Ȳ::Matrix, T, nrows, ncols)
+function _densification_pullback(Ȳ::AbstractMatrix, T, nrows, ncols)
  row_idxs = cumsum(nrows) .- nrows .+ 1
  col_idxs = cumsum(ncols) .- ncols .+ 1
 

diff --git a/test/chainrules.jl b/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