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

subset and merge for VarInfo (clean version) #544

Merged
merged 28 commits into from
Oct 19, 2023
Merged

Conversation

torfjelde
Copy link
Member

@torfjelde torfjelde commented Oct 9, 2023

I was thinking that #542 would be a quick merge, and so I just based #543 on #542 so I could easily test the new Gibbs in Turing.jl, buuut it seems that wasn't as quick as I thought and so this PR replaces #543 without any dependence on #542.

Hence: see #543 for examples and some comments

@torfjelde torfjelde changed the title Added subset and merge for VarInfo subset and merge for VarInfo (clean version) Oct 9, 2023
src/varinfo.jl Outdated
"""
merge(varinfo_left::VarInfo, varinfo_right::VarInfo)

Merge two `VarInfo` instances into one, giving precedence to `varinfo_right` when reasonable.
Copy link
Member Author

Choose a reason for hiding this comment

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

I will add some more docs for this a bit later today

@codecov
Copy link

codecov bot commented Oct 9, 2023

Codecov Report

Attention: 24 lines in your changes are missing coverage. Please review.

Comparison is base (927799f) 81.51% compared to head (d3a9b56) 82.76%.

Additional details and impacted files
@@            Coverage Diff             @@
##           master     #544      +/-   ##
==========================================
+ Coverage   81.51%   82.76%   +1.24%     
==========================================
  Files          25       25              
  Lines        3003     3180     +177     
==========================================
+ Hits         2448     2632     +184     
+ Misses        555      548       -7     
Files Coverage Δ
src/DynamicPPL.jl 33.33% <ø> (ø)
src/test_utils.jl 86.53% <100.00%> (+0.10%) ⬆️
src/simple_varinfo.jl 75.25% <76.47%> (+3.98%) ⬆️
src/varinfo.jl 93.31% <96.18%> (+0.53%) ⬆️
src/abstract_varinfo.jl 85.84% <36.36%> (-5.34%) ⬇️
src/threadsafe.jl 45.61% <60.00%> (+25.20%) ⬆️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@torfjelde
Copy link
Member Author

Do any of you want to take a look at this @devmotion @yebai ?

@yebai yebai requested a review from sunxd3 October 10, 2023 12:16
@yebai
Copy link
Member

yebai commented Oct 10, 2023

I'll take a look later today. @sunxd3 can you also help review this PR?

src/varinfo.jl Outdated Show resolved Hide resolved
end

function subset(metadata::Metadata, vns::AbstractVector{<:VarName})
# TODO: Should we error if `vns` contains a variable that is not in `metadata`?
Copy link
Member

Choose a reason for hiding this comment

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

At least a warning?

src/varinfo.jl Outdated
metadata = merge_metadata(varinfo_left.metadata, varinfo_right.metadata)
lp = getlogp(varinfo_left) + getlogp(varinfo_right)
# TODO: Is this really the way we want to combine `num_produce`?
num_produce = varinfo_left.num_produce[] + varinfo_right.num_produce[]
Copy link
Member

Choose a reason for hiding this comment

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

Remind me the use of num_produce?

Copy link
Member

Choose a reason for hiding this comment

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

Semantically, taking the sum of two num_produce produces no meaning. I think we should throw an error if both num_produce are non-zero.

Copy link
Member Author

@torfjelde torfjelde Oct 12, 2023

Choose a reason for hiding this comment

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

num_produce is used by the particle samplers to indicate how many observations we've seen / what's the current "observation index".

Semantically, taking the sum of two num_produce produces no meaning. I think we should throw an error if both num_produce are non-zero.

I don't think that's quite right. It makes sense if you go "execute model A and then separately execute model B, and now we want to merge the two"; to ensure that the observation index is correctly, we need to add them, no?

Is this a scenario we want to consider? Probably not. But to me it wasn't obvious what else to do.

No matter, I don't think erroring if num_produce is non-zero is the right way to go. We should allow something like

varinfo_with_num_produce = last(evaluate!!(model, varinfo, context))
merge(varinfo_with_num_produce, varinfo)

If there's no meaning in adding them, then I suggest the alternative is to just give precedence to varinfo_right, as is the semantics of merge (if a field/property/key is present in both left and right, then the value in right takes precedence).

EDIT: It's fair to ask if we should also do that for logp. I went with sum because I was imagining scenarios where we want to do something like execute two (sub-)models in parallel and then merge the resulting varinfos. In such a scenario, we want to add the logp fields together. But we could maybe make this a separate function, e.g. merge_with_add_logp or something (preferably named in a better way).

Copy link
Member Author

Choose a reason for hiding this comment

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

So what should we do here? Give precedence to varinfo_right or leave as it is?

@sunxd3
Copy link
Member

sunxd3 commented Oct 10, 2023

Looks good to me, but I am not very familiar with VarInfo. Left some questions mainly for my understanding.

test/varinfo.jl Outdated Show resolved Hide resolved
Copy link
Member

@yebai yebai left a comment

Choose a reason for hiding this comment

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

It looks like a lot of additional code to support VarInfo. Is there a plan to drop VarInfo in favour of SimpleVarInfo at some point?

We have nearly everything we need in SimpleVarInfo now. So, the marginal benefit of keeping both VarInfo and SimpleVarInfo is diminishing.

src/varinfo.jl Outdated
metadata = merge_metadata(varinfo_left.metadata, varinfo_right.metadata)
lp = getlogp(varinfo_left) + getlogp(varinfo_right)
# TODO: Is this really the way we want to combine `num_produce`?
num_produce = varinfo_left.num_produce[] + varinfo_right.num_produce[]
Copy link
Member

Choose a reason for hiding this comment

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

Semantically, taking the sum of two num_produce produces no meaning. I think we should throw an error if both num_produce are non-zero.

@@ -1331,6 +1669,15 @@ function setorder!(vi::VarInfo, vn::VarName, index::Int)
return vi
end

"""
getorder(vi::VarInfo, vn::VarName)
Copy link
Member

Choose a reason for hiding this comment

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

I like this API -- we can consider depreciating num_produce in favour of getorder in the longer run.

Copy link
Member Author

Choose a reason for hiding this comment

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

I'm confused. Isn't num_produce and order different things? At least they are two different fields in VarInfo. I I just added a getorder method because I've been trying to make the interface of VarInfo simpler.

Copy link
Member

Choose a reason for hiding this comment

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

num_produce is used to track the current observation index (for all vns), and its current value inserted to VarInfo.metadata when a new vn is created.

@torfjelde
Copy link
Member Author

We have nearly everything we need in SimpleVarInfo now. So, the marginal benefit of keeping both VarInfo and SimpleVarInfo is diminishing.

We don't though 😕 VarInfo is much more flexible than SimpleVarInfo is. See #528, and defining such a "vector" (which is going to be very simliar to metadata) is going to be challenging to get right.

Until that has been done, we need to continue supporting VarInfo.

@yebai
Copy link
Member

yebai commented Oct 12, 2023

See #528, and defining such a "vector" (which is going to be very simliar to metadata) is going to be challenging to get right.

Yes, I liked the VarDict abstraction as a replacement for Metadata. After this VarDict is complete, SimpleVarInfo should be more general and modular than VarInfo.

@torfjelde
Copy link
Member Author

Yes, I liked the VarDict abstraction as a replacement for Metadata. After this VarDict is complete, SimpleVarInfo should be more general and modular than VarInfo.

Agreed, but that will take time; it's non-trivial to get right IMO.

# the `eltype` to `VarName`? This might be useful when someone does `[@varname(x[1]), @varname(m)]` which
# might result in a `Vector{Any}`.
"""
subset(varinfo::AbstractVarInfo, vns::AbstractVector{<:VarName})
Copy link
Member

Choose a reason for hiding this comment

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

If we would not already have so many getindex methods, I would have thought that getindex would be a natural name for this function. But maybe it's still an option?

Then we could have getindex(::AbstractVarInfo, ::AbstractVector{<:VarName}) -> AbstractVarInfo and getindex(::T, ::VarName) -> typeof_varname_variate, similar to [1,2,3][[1,3]] = [1, 3] and [1,2,3][2] = 2.

Copy link
Member Author

Choose a reason for hiding this comment

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

I'd really like this yes, but I also really don't want to touch getindex in this codebase 😅

Happy to make this a long-term goal or something though!

See docstring of [`subset(varinfo, vns)`](@ref) for examples.
"""
function Base.merge(varinfo::AbstractVarInfo, varinfo_others::AbstractVarInfo...)
return merge(Base.merge(varinfo, first(varinfo_others)), Base.tail(varinfo_others)...)
Copy link
Member

Choose a reason for hiding this comment

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

It's possible that length(varinfo_others) == 0, and then the function would error it seems?

Copy link
Member Author

Choose a reason for hiding this comment

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

Ah true; I'll address this.

I meant to add tests for merge with more than two inputs too, but didn't get around to it. Will do that too.

src/abstract_varinfo.jl Outdated Show resolved Hide resolved
src/simple_varinfo.jl Outdated Show resolved Hide resolved
@torfjelde
Copy link
Member Author

Okay, so I think this is now properly ready for a potential approval (after comments have been resolved)

@coveralls
Copy link

coveralls commented Oct 13, 2023

Pull Request Test Coverage Report for Build 6518813946

  • 160 of 184 (86.96%) changed or added relevant lines in 5 files are covered.
  • 1 unchanged line in 1 file lost coverage.
  • Overall coverage increased (+1.2%) to 82.767%

Changes Missing Coverage Covered Lines Changed/Added Lines %
src/simple_varinfo.jl 13 17 76.47%
src/varinfo.jl 126 131 96.18%
src/abstract_varinfo.jl 4 11 36.36%
src/threadsafe.jl 12 20 60.0%
Files with Coverage Reduction New Missed Lines %
src/varinfo.jl 1 93.32%
Totals Coverage Status
Change from base Build 6477265825: 1.2%
Covered Lines: 2632
Relevant Lines: 3180

💛 - Coveralls

@torfjelde
Copy link
Member Author

Okay, so I added quite a bit of changes to ensure that ThreadSafeVarInfo will also work with all of this. Atm supporting these functionalities for ThreadSafeVarInfo is not really necessary unless we start using merge and subset within models (which I do imagine we'll start doing at some point, as it allows us to easily parallelize even sampling!), but it's always annoying when ThreadSafeVarInfo doesn't actually_ implement the varinfo interface so figured we should add it.

@torfjelde
Copy link
Member Author

Test coverage should now be improved 👍

@torfjelde
Copy link
Member Author

This is ready for a look-over now btw.

Copy link
Member

@yebai yebai left a comment

Choose a reason for hiding this comment

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

Thanks, @torfjelde -- it looks good to me!

@yebai yebai enabled auto-merge October 19, 2023 13:26
@yebai yebai added this pull request to the merge queue Oct 19, 2023
Merged via the queue into master with commit efd9da3 Oct 19, 2023
12 of 13 checks passed
@yebai yebai deleted the torfjelde/subset-and-merge branch October 19, 2023 16:48
@torfjelde
Copy link
Member Author

@devmotion had some open comments here, so I really would have preferred to have his 👍 on this PR too

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.

5 participants