RFC: Allows blocks to be <:AbstractVector or <:Tuple#56
RFC: Allows blocks to be <:AbstractVector or <:Tuple#56nickrobinson251 wants to merge 5 commits into
Conversation
Codecov Report
@@ Coverage Diff @@
## master #56 +/- ##
==========================================
- Coverage 94.04% 90.83% -3.21%
==========================================
Files 5 5
Lines 235 240 +5
==========================================
- Hits 221 218 -3
- Misses 14 22 +8
Continue to review full report at Codecov.
|
I for one definitely want this. |
|
Julia v1.0 is grumpy. I think the fix is to add |
5477835 to
70ddc17
Compare
| blocks3 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)] | ||
|
|
||
| @testset for V in (Tuple, Vector) | ||
| b1 = BlockDiagonal(V(blocks1)) |
There was a problem hiding this comment.
we shold introduce an indent under this new @testset for V in (Tuple, Vector) loop, but that'd give us huge diffs, and make it harder to see what's changed in the tests, so to make review easier i'll make that aesthetic change in a follow-up PR
| bm::BlockDiagonal{T, V}, | ||
| v::StridedVector{T} | ||
| ) where {T<:Union{Real, Complex}, V<:Matrix{T}} | ||
| ) where {T<:Union{Real, Complex}, V} |
There was a problem hiding this comment.
We should continue to maintain strict type constraints here as this rule is not in general correct for blocks that aren't Matrix{<:Real}s (or nearly-Matrix{<:Real}s, like StridedMatrix{<:Real}s). Something like
V<:Vector{<:Matrix{<:Real}}would actually be a good choice to my mind. See e.g. JuliaDiff/ChainRules.jl#232
There was a problem hiding this comment.
love it. makes this PR much easier 😂
- This is a breaking change, as the type param `V` is now different,
being the type of the `blocks` collection, rather than the `eltype`
of that collection.
- This allows allows the collection to be something other than a
`Vector` e.g.
- a `CuArray` of matrices (allowing use on GPU hopefully)
- a `Tuple` with different concrete subtypes of `AbstractMatrix{T}`
(avoiding heterogeneous block to being abstractly typed
i.e. a way to avoid `Vector{AbstractMatrix{T}}`)
70ddc17 to
c9bf915
Compare
| @@ -1,5 +1,5 @@ | |||
| # constructor | |||
| function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where {V} | |||
| function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::V) where {V<:AbstractVector} | |||
There was a problem hiding this comment.
i think AbstractVector is fine here, but i'll rely on @willtebbutt to tell me if not
There was a problem hiding this comment.
If we exclude the tuple we won't have an rrule for constructors of BlockDiagonals with tuple blocks. Are you asking if it's ok not to have rrule for tuple blocks, or am I missing the point of the question?
There was a problem hiding this comment.
we won't have an rrule for constructors of BlockDiagonals with tuple blocks
correct, that'd be a nice feature to add as a follow-up, but out-of-scope for this PR
But my question is a different one. Before this PR, blocks were always a Vector, so we had rrule(::Type{<:BlockDiagonal}, blocks::Vector) whereas this PR allows them to be any AbstractVector (or Tuple but ignore that, since we're not adding that rrule here). So my question is, now that blocks can be any AbstractVector, are we okay with broadening the rrule from rrule(::Type{<:BlockDiagonal}, blocks::Vector) to rrule(::Type{<:BlockDiagonal}, blocks::AbstractVector) ?
(it's a question about how tightly to constrain rrules #56 (comment))
| bm::BlockDiagonal{T, V}, | ||
| v::StridedVector{T} | ||
| ) where {T<:Union{Real, Complex}, V<:Matrix{T}} | ||
| ) where {T<:Union{Real, Complex}, V<:Vector{Matrix{T}}} |
9c0b6d3 to
d5d29da
Compare
| end | ||
| !!! info "`V` type" | ||
| `blocks::V` should be a `Tuple` or `AbstractVector` where each component (each block) is | ||
| `<:AbstractMatrix{T}` for some common element type `T`. |
There was a problem hiding this comment.
do we want to constrain T to at least a Number, or are we happy with anything?
There was a problem hiding this comment.
i think if T isn't Number-like, some methods will fail. But 🤷 Would there be any gain from constraining things?
| # BlockArrays-like functions | ||
| """ | ||
| blocksizes(B::BlockDiagonal) -> Vector{Tuple} | ||
| blocksizes(B::BlockDiagonal{T, V}) -> V |
There was a problem hiding this comment.
Do we want this to return a Vector or a Tuple depending on the block type? Or should we force it to return a vector/tuple always?
There was a problem hiding this comment.
my thinking is that we'd use blocks::Tuple for performance reasons when you've a small number of heterogenous blocks, in which case getting small Tuples from methods that operate on blocks will also be slightly better for performance, so no reason to force a Vector to be returned. It doesn't make sense to force a Tuple return in general, since there could be many elements.
| @@ -1,5 +1,5 @@ | |||
| # constructor | |||
| function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::Vector{V}) where {V} | |||
| function ChainRulesCore.rrule(::Type{<:BlockDiagonal}, blocks::V) where {V<:AbstractVector} | |||
There was a problem hiding this comment.
If we exclude the tuple we won't have an rrule for constructors of BlockDiagonals with tuple blocks. Are you asking if it's ok not to have rrule for tuple blocks, or am I missing the point of the question?
| U_Vt = ntuple(length(blocks(B))) do i | ||
| F = svd(getblock(B, i), full=full) | ||
| append!(S, F.S) | ||
| (F.U, F.Vt) |
There was a problem hiding this comment.
How come we aren't push!ing like above? Or vice versa if first.()/last.() is better?
There was a problem hiding this comment.
because above we want U (likewise Vt) to be a Vector, so create a Vector and add the elements. Whereas here we want it to be a Tuple hence ntuple, and first.(ntup) will still give us a tuple.
rather than push! we should probably be preallocating a Vector of the right length and overwriting the values, but the two functions would still be different (since Tuples are immutable)
|
|
||
| function BlockDiagonal(blocks::Vector{V}) where {T, V<:AbstractMatrix{T}} | ||
| function BlockDiagonal(blocks::V) where { | ||
| T, V<:Union{Tuple{Vararg{<:AbstractMatrix{T}}}, AbstractVector{<:AbstractMatrix{T}}} |
There was a problem hiding this comment.
Do you really want all block matrices to have the same eltype? Or could you live with them having different eltypes, and T being promote_type(map(eltype, blocks)...)? See LinearMaps.jl. If you need them to be homogeneous, then you may wish to include some promotion mechanism.
There was a problem hiding this comment.
could you live with them having different eltypes, and
Tbeingpromote_type(map(eltype, blocks)...)?
that's a great suggestion
to be honest, i hadn't considered it. This currently just keeps T with the same meaning as on master.
Are there any downsides to T being promote_type(map(eltype, blocks)...)?
There was a problem hiding this comment.
trying this locally, and running tests, it does cause some inference issues for rrule
BlockDiagonal * Vector: Error During Test at /Users/npr/repos/BlockDiagonals.jl/test/chainrules.jl:12
Got exception outside of a @test
return type Array{Float64,1} does not match inferred return type Any
Stacktrace:
[1] error(::String) at ./error.jl:33
[2] _test_inferred(::Function, ::InplaceableThunk{Thunk{BlockDiagonals.var"#13#19"{Array{Float64,1},BlockDiagonal{Float64,Array{Array{Float64,2},1}}}},BlockDiagonals.var"#14#20"{Array{Float64,1},BlockDiagonal{Float64,Array{Array{Float64,2},1}}}}; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/npr/.julia/packages/ChainRulesTestUtils/2QRer/src/testers.jl:228
[3] _test_inferred(::Function, ::InplaceableThunk{Thunk{BlockDiagonals.var"#13#19"{Array{Float64,1},BlockDiagonal{Float64,Array{Array{Float64,2},1}}}},BlockDiagonals.var"#14#20"{Array{Float64,1},BlockDiagonal{Float64,Array{Array{Float64,2},1}}}}) at /Users/npr/.julia/packages/ChainRulesTestUtils/2QRer/src/testers.jl:227
[4] test_rrule(::typeof(*), ::BlockDiagonal{Float64,Array{Array{Float64,2},1}}, ::Vararg{Any,N} where N; output_tangent::ChainRulesTestUtils.Auto, fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5,1,FiniteDifferences.UnadaptedFiniteDifferenceMethod{7,5}}, check_inferred::Bool, fkwargs::NamedTuple{(),Tuple{}}, rtol::Float64, atol::Float64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/npr/.julia/packages/ChainRulesTestUtils/2QRer/src/testers.jl:187
[5] test_rrule(::Function, ::BlockDiagonal{Float64,Array{Array{Float64,2},1}}, ::Vararg{Any,N} where N) at /Users/npr/.julia/packages/ChainRulesTestUtils/2QRer/src/testers.jl:153
[6] top-level scope at /Users/npr/repos/BlockDiagonals.jl/test/chainrules.jl:15
perhaps someone more familiar with ChainRules has some insight on why this might be (@willtebbutt , @mzgubic)?
There was a problem hiding this comment.
I don't know why this is (I assume the answer contains the word compiler). But I'd rather live with the same eltype than with this I think
There was a problem hiding this comment.
Okay, this change would be non-breaking anyway, so let's leave it as a potentially follow-up #62
There was a problem hiding this comment.
I think that's perfectly fine. But then you may wish to consider to include a constructor that by itself does not yet require equal eltypes, but first determines T via promotion and then map(B -> convert(AbstractMatrix{T}, B), blocks). This should be a no-op if all matrices have the same eltype, but otherwise relieve you from annoying method errors. Just a suggestion.
Vis now different,being the type of the
blockscollection, rather than theeltypeof that collection.
Vectore.g.CuArrayof matrices (allowing use on GPU hopefully)Tuplewith different concrete subtypes ofAbstractMatrix{T}(avoiding heterogeneous block to being abstractly typed
i.e. a way to avoid
Vector{AbstractMatrix{T}})TODO:
right now this just does the minimal (ugly) thing of looping over existing tests with bothTupleandVectorblocks -- we should probably refactor tests to make this nicer, and make future tests easier to add... probably in this PR if we can avoid making the diffs to big.