Test accepting thunks passed to pullbacks (#73)

Author: mzgubic

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.18"
+version = "0.1.19"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -11,7 +11,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
 ChainRulesCore = "0.9.44, 0.10"
-ChainRulesTestUtils = "0.6.3, 0.7"
+ChainRulesTestUtils = "0.7.10"
 FillArrays = "0.6, 0.7, 0.8, 0.9, 0.10, 0.11"
 FiniteDifferences = "0.12.3"
 julia = "1"

src/chainrules.jl

@@ -1,25 +1,30 @@
 # constructor
+_BlockDiagonal_pullback(Δ::Tangent) = (NoTangent(), Δ.blocks)
+_BlockDiagonal_pullback(Δ::AbstractThunk) = _BlockDiagonal_pullback(unthunk(Δ))
 function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where {V}
-    BlockDiagonal_pullback(Δ::Tangent) = (NoTangent(), Δ.blocks)
-    return BlockDiagonal(blocks), BlockDiagonal_pullback
+    return BlockDiagonal(blocks), _BlockDiagonal_pullback
 end
 
 # densification
+function _densification_pullback(Ȳ::Matrix, T, nrows, ncols)
+    row_idxs = cumsum(nrows) .- nrows .+ 1
+    col_idxs = cumsum(ncols) .- ncols .+ 1
+
+    Δblocks = map(eachindex(nrows)) do n
+        block_rows = row_idxs[n]:(row_idxs[n] + nrows[n] - 1)
+        block_cols = col_idxs[n]:(col_idxs[n] + ncols[n] - 1)
+        return Ȳ[block_rows, block_cols]
+    end
+    return (NoTangent(), Tangent{T}(blocks=Δblocks))
+end
+function _densification_pullback(Ȳ::AbstractThunk, T, nrows, ncols)
+    return _densification_pullback(unthunk(Ȳ), T, nrows, ncols)
+end
 function ChainRulesCore.rrule(::Type{<:Base.Matrix}, B::T) where {T<:BlockDiagonal}
     nrows = size.(B.blocks, 1)
     ncols = size.(B.blocks, 2)
-    function Matrix_pullback(Δ::Matrix)
-        row_idxs = cumsum(nrows) .- nrows .+ 1
-        col_idxs = cumsum(ncols) .- ncols .+ 1
-
-        Δblocks = map(eachindex(nrows)) do n
-            block_rows = row_idxs[n]:(row_idxs[n] + nrows[n] - 1)
-            block_cols = col_idxs[n]:(col_idxs[n] + ncols[n] - 1)
-            return Δ[block_rows, block_cols]
-        end
-        return (NoTangent(), Tangent{T}(blocks=Δblocks))
-    end
-    return Matrix(B), Matrix_pullback
+    densification_pullback(ȳ) = _densification_pullback(ȳ, T, nrows, ncols)
+    return Matrix(B), densification_pullback
 end
 
 # multiplication
@@ -37,23 +42,20 @@ function ChainRulesCore.rrule(
     row_idxs = cumsum(nrows) .- nrows .+ 1
     col_idxs = cumsum(ncols) .- ncols .+ 1
 
-    function bm_vector_mul_pullback(Δ)
+    function bm_vector_mul_pullback(Δy)
+        ȳ = unthunk(Δy)
         Δ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)
             return InplaceableThunk(
-                @thunk(Δ[block_rows] * v[block_cols]'),
-                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
-        return (
-            NoTangent(),
-            Tangent{BlockDiagonal{T, V}}(;blocks=Δblocks),
-            InplaceableThunk(
-                @thunk(bm' * Δ),
-                X̄ -> mul!(X̄, bm', Δ, true, true)
-            ),
-        )
+
+        b̄m = Tangent{BlockDiagonal{T, V}}(;blocks=Δblocks)
+        v̄ = InplaceableThunk(@thunk(bm' * ȳ), X̄ -> mul!(X̄, bm', ȳ, true, true))
+        return NoTangent(), b̄m, v̄
     end
     return y, bm_vector_mul_pullback
 end

test/runtests.jl

@@ -6,6 +6,8 @@ using FiniteDifferences # For overloading to_vec
 using Test
 using LinearAlgebra
 
+push!(ChainRulesTestUtils.TRANSFORMS_TO_ALT_TANGENTS, x -> @thunk(x))
+
 @testset "BlockDiagonals" begin
     # The doctests fail version other than 64bit julia 1.6.x, due to printing differences
     Sys.WORD_SIZE == 64 && v"1.6" <= VERSION < v"1.7" && doctest(BlockDiagonals)