Merge pull request #77 from invenia/mz/cr1

Update to ChainRules 1.0

Author: mzgubic

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.20"
+version = "0.1.21"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -10,8 +10,8 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
-ChainRulesCore = "0.9.44, 0.10"
-ChainRulesTestUtils = "0.7.10"
+ChainRulesCore = "1"
+ChainRulesTestUtils = "1"
 FillArrays = "0.6, 0.7, 0.8, 0.9, 0.10, 0.11, 0.12"
 FiniteDifferences = "0.12.3"
 julia = "1"

src/BlockDiagonals.jl

@@ -6,6 +6,8 @@ using FillArrays: Zeros
 using FiniteDifferences
 using LinearAlgebra
 
+import ChainRulesCore.ProjectTo
+
 export BlockDiagonal, blocks
 export blocksize, blocksizes, nblocks
 

src/chainrules.jl

@@ -1,6 +1,7 @@
 # constructor
 _BlockDiagonal_pullback(Δ::Tangent) = (NoTangent(), Δ.blocks)
 _BlockDiagonal_pullback(Δ::AbstractThunk) = _BlockDiagonal_pullback(unthunk(Δ))
+_BlockDiagonal_pullback(Δ::BlockDiagonal) = (NoTangent(), Δ.blocks)
 function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where {V}
     return BlockDiagonal(blocks), _BlockDiagonal_pullback
 end
@@ -48,14 +49,34 @@ function ChainRulesCore.rrule(
             block_rows = row_idxs[i]:(row_idxs[i] + nrows[i] - 1)
             block_cols = col_idxs[i]:(col_idxs[i] + ncols[i] - 1)
             return InplaceableThunk(
+                X̄ -> mul!(X̄, ȳ[block_rows], v[block_cols]', true, true),
                 @thunk(ȳ[block_rows] * v[block_cols]'),
-                X̄ -> mul!(X̄, ȳ[block_rows], v[block_cols]', true, true)
             )
         end
 
         b̄m = Tangent{BlockDiagonal{T, V}}(;blocks=Δblocks)
-        v̄ = InplaceableThunk(@thunk(bm' * ȳ), X̄ -> mul!(X̄, bm', ȳ, true, true))
+        v̄ = InplaceableThunk(X̄ -> mul!(X̄, bm', ȳ, true, true), @thunk(bm' * ȳ))
         return NoTangent(), b̄m, v̄
     end
     return y, bm_vector_mul_pullback
 end
+
+function ProjectTo(b::BlockDiagonal)
+    blocks = map(ProjectTo, b.blocks)
+    return ProjectTo{BlockDiagonal}(; blocks=blocks, blocksizes=blocksizes(b))
+end
+
+function (project::ProjectTo{BlockDiagonal})(dx)
+    # prepare to index into the dense array
+    nrows = first.(project.blocksizes)
+    ncols = last.(project.blocksizes)
+    row_idxs = cumsum(nrows) .- nrows .+ 1
+    col_idxs = cumsum(ncols) .- ncols .+ 1
+    # project each block individually
+    blocks = map(eachindex(nrows)) do i
+        block_rows = row_idxs[i]:(row_idxs[i] + nrows[i] - 1)
+        block_cols = col_idxs[i]:(col_idxs[i] + ncols[i] - 1)
+        project.blocks[i](dx[block_rows, block_cols])
+    end
+    return BlockDiagonal(blocks)
+end

test/chainrules.jl

@@ -2,6 +2,7 @@
     @testset "BlockDiagonal" begin
         x = [randn(1, 2), randn(2, 2)]
         test_rrule(BlockDiagonal, x)
+        test_rrule(BlockDiagonal, x; output_tangent=Tangent{BlockDiagonal}(;blocks=x))
     end
 
     @testset "Matrix" begin
@@ -14,4 +15,13 @@
         v = rand(6)
         test_rrule(*, D, v)
     end
+
+    @testset "ProjectTo" begin
+        bd = BlockDiagonal([ones(2, 2), ones(3, 3)])
+        project = ProjectTo(bd)
+        @test project(ones(5, 5)) == bd
+        @test project(adjoint(ones(5, 5))) == bd
+        @test project(Diagonal(ones(5))) isa BlockDiagonal
+        @test project(Diagonal(ones(5))) == Diagonal(ones(5))
+    end
 end