Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RFC: Allows blocks to be <:AbstractVector or <:Tuple #56

Closed
wants to merge 5 commits into from

Conversation

nickrobinson251
Copy link
Contributor

@nickrobinson251 nickrobinson251 commented Feb 2, 2021

TODO:

  • decide if we want this. I think we do... but i'd be interested in more feedback!
  • make tests nicer
    • right now this just does the minimal (ugly) thing of looping over existing tests with both Tuple and Vector blocks -- 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.
    • open a a follow-up PR with the aesthetic changes
  • add tests with different collection types (e.g. CuArrays or NamedDimsArrays)
  • maybe add docs (doctests?) showing the different options / trade-offs for using Tuples versus Vectors?
  • [optional] benchmarks would be interesting (e.g. what's best when you have 20 non-homogenous blocks?) but may be best in a follow up PR (also we just want more benchmarks in general Add benchmarks #19)

@codecov
Copy link

codecov bot commented Feb 2, 2021

Codecov Report

Merging #56 (d79d842) into master (fda7259) will decrease coverage by 3.20%.
The diff coverage is 72.72%.

Impacted file tree graph

@@            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     
Impacted Files Coverage Δ
src/chainrules.jl 100.00% <ø> (ø)
src/linalg.jl 82.85% <62.50%> (-10.80%) ⬇️
src/blockdiagonal.jl 82.14% <100.00%> (+0.32%) ⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fda7259...d5d29da. Read the comment docs.

@willtebbutt
Copy link
Contributor

decide if we want this. I think we do... but i'd be interested in more feedback!

I for one definitely want this.

@nickrobinson251
Copy link
Contributor Author

Julia v1.0 is grumpy. I think the fix is to add cumsum(itr) to Compat.jl JuliaLang/Compat.jl#737

test/linalg.jl Outdated Show resolved Hide resolved
@nickrobinson251 nickrobinson251 force-pushed the npr/allow-tuple-or-any-vec branch 2 times, most recently from 5477835 to 70ddc17 Compare February 15, 2021 21:53
blocks3 = [rand(rng, N1, N1), rand(rng, N2, N2), rand(rng, N2, N2)]

@testset for V in (Tuple, Vector)
b1 = BlockDiagonal(V(blocks1))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

test/chainrules.jl Outdated Show resolved Hide resolved
test/chainrules.jl Outdated Show resolved Hide resolved
@@ -27,7 +27,7 @@ function ChainRulesCore.rrule(
::typeof(*),
bm::BlockDiagonal{T, V},
v::StridedVector{T}
) where {T<:Union{Real, Complex}, V<:Matrix{T}}
) where {T<:Union{Real, Complex}, V}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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}}`)
@@ -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}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think AbstractVector is fine here, but i'll rely on @willtebbutt to tell me if not

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

@nickrobinson251 nickrobinson251 Feb 16, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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))

@@ -27,7 +27,7 @@ function ChainRulesCore.rrule(
::typeof(*),
bm::BlockDiagonal{T, V},
v::StridedVector{T}
) where {T<:Union{Real, Complex}, V<:Matrix{T}}
) where {T<:Union{Real, Complex}, V<:Vector{Matrix{T}}}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this okay, @willtebbutt ?

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`.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we want to constrain T to at least a Number, or are we happy with anything?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think if T isn't Number-like, some methods will fail. But 🤷 Would there be any gain from constraining things?


Return the on-diagonal blocks of B.
"""
blocks(B::BlockDiagonal) = B.blocks

# BlockArrays-like functions
"""
blocksizes(B::BlockDiagonal) -> Vector{Tuple}
blocksizes(B::BlockDiagonal{T, V}) -> V
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How come we aren't push!ing like above? Or vice versa if first.()/last.() is better?

Copy link
Contributor Author

@nickrobinson251 nickrobinson251 Feb 16, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

end

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}}}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could you live with them having different eltypes, and T being promote_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)...)?

Copy link
Contributor Author

@nickrobinson251 nickrobinson251 Feb 16, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, this change would be non-breaking anyway, so let's leave it as a potentially follow-up #62

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Turn vector of matrices into a tuple?
4 participants