Makes _setval! and thus the prob-macro compatible with v4

[torfjelde]

Currently the following fails:

DynamicPPL._setval!(vi, c)

where vi is a VarInfo and c is a MCMCChains.Chains if the model contains non-univariate variables. This PR provides a fix by using the new MCMCChains.namesingroup to get the corresponding index-symbols which is then used to extract the values from the chain c.

The issue is that after v4.0 vi.metadata[key].vns results in a VarName which is "identical" to key in the sense that even if a represents a multivariate variable, vi.metadata[:a].vns == [VarName(:a, ())], not the varnames corresponding to the indices for that variable. See TuringLang/Turing.jl#1352 for a related issue.

[cpfiffer]

I'll provisionally approve, but since DynamicPPL isn't really my bag I'd like to see a merge in from one of the "real" DynamicPPL people.

[yebai]

Ready to merge once the tests pass.

[yebai]

Can we copy-and-paste the PR comments here, so we can remind ourselves what's happening in the future?

[torfjelde]

There's an issue; I forgot that DynamicPPL.jl doesn't depend on MCMCChains.jl but only on AbstractMCMC.jl (which is where we get the AbstractChains from). Essentially, now that .vns doesn't give you the elementwise indices, the old implementation is broken. But we also cannot use the group from MCMCChains.jl to "correctly" extract the wanted elements because DynamicPPL.jl doesn't depend on MCMCChains.jl.

So either we

1. revert the changes made to vns so that indeed vns gives you the indices of all the elements,
2. make the group method (or some other fix) part of AbstractMCMC.jl.

Personally I strongly prefer option (1). At the moment vns seems completely redundant since it just replicates the key already used to extract this particular variable, e.g. vi.metadata[:a].vns == [VarName(:a, ())] is true.

I'm also a bit out of the loop here, but it's a bit strange to me if we implement methods for AbstractChains from AbstractMCMC.jl, but we end up assuming we have the indexing behavior of MCMCChains.jl. Seems to me that either we depend directly on MCMCChains.jl, or we need to describe some desired indexing behavior that AbstractChains should have and then tell people to implement that if they're implementing some "chains" type. Though, it's also not clear to me how many different AbstractChains sub-types we're going to see in the wild. In addition, the only place where AbstractChains is actually used in the entirety of DynamicPPL.jl is here:

DynamicPPL.jl/src/prob_macro.jl

Line 231 in f670195

_setval!(vi::TypedVarInfo, c::AbstractChains) = _setval!(vi.metadata, vi, c)

. Seems a bit overkill, no?

@mohamed82008 @devmotion Maybe you guys have thought about this before?

[cpfiffer]

So either we

1. revert the changes made to `vns` so that indeed `vns` gives you the indices of all the elements,

2. make the `group` method (or some other fix) part of AbstractMCMC.jl.

I don't think 2 is a good option -- group is specific to MCMCChains and I can't yet see a convincing case why it should be added upstream.

So I prefer option 1 here.

I'm also a bit out of the loop here, but it's a bit strange to me if we implement methods for AbstractChains from AbstractMCMC.jl, but we end up assuming we have the indexing behavior of MCMCChains.jl. Seems to me that either we depend directly on MCMCChains.jl, or we need to describe some desired indexing behavior that AbstractChains should have and then tell people to implement that if they're implementing some "chains" type. Though, it's also not clear to me how many different AbstractChains sub-types we're going to see in the wild. In addition, the only place where AbstractChains is actually used in the entirety of DynamicPPL.jl is here:

I agree that it's strange. So far, we haven't moved any of the indexing functionality upstream to AbstractMCMC, and it's not clear that we would want to -- I think we are implicitly requiring getindex from chains without formalizing the specification. Regardless, there aren't currently any chain competitors that implement the methods, so it's kind of a moot point at the moment.

[torfjelde]

I don't think 2 is a good option -- group is specific to MCMCChains and I can't yet see a convincing case why it should be added upstream.

Fair point 👍

Regardless, there aren't currently any chain competitors that implement the methods, so it's kind of a moot point at the moment.

I wouldn't say it's moot though; AbstractChains conveys information to the user that "if you implement this, then it works with anything that requires AbstractChains!", which is not the case, e.g. _setval! requires AbstractChains but will fail if you're using any other implementation than Chains. Then, imo, it should just directly rely on Chains instead.

I'm fully aware that AbstractChains likely was implemented as a quick-fix to deal with dependency-issues or something, but I'm just saying that I wouldn't say it's a moot point as it inevitably will lead to a lot of frustration for someone down the road (e.g. me 😢 ).

[cpfiffer]

I'm fully aware that AbstractChains likely was implemented as a quick-fix to deal with dependency-issues or something, but I'm just saying that I wouldn't say it's a moot point as it inevitably will lead to a lot of frustration for someone down the road (e.g. me cry

It's not necessarily intended to be a quick fix, I think my perception of why it's a moot point is that we should only really put stuff in AbstractMCMC if there's a realistic default function to call, which is not the case for an arbitrary chain type. At least, we haven't formalized things to allow a convenient default.

Honestly this should probably be Chains instead of AbstractChains to be more specific, and all the code that touches MCMCChains should be moved behind an @requires block.

[torfjelde]

It's not necessarily intended to be a quick fix, I think my perception of why it's a moot point is that we should only really put stuff in AbstractMCMC if there's a realistic default function to call, which is not the case for an arbitrary chain type. At least, we haven't formalized things to allow a convenient default.

Yeah, I think we agree here: it doesn't make sense to have an abstract type for which we define behavior for if we do not also define what's expected behavior for this type.

Honestly this should probably be Chains instead of AbstractChains to be more specific, and all the code that touches MCMCChains should be moved behind an @requires block.

Not sure if @requires is the right approach though as this functionality is necessary for the probability-macros for a model. Maybe all this should just go into Turing.jl rather than DynamicPPL.jl?

[cpfiffer]

Not sure if @requires is the right approach though as this functionality is necessary for the probability-macros for a model. Maybe all this should just go into Turing.jl rather than DynamicPPL.jl?

Maybe? I dunno. I could see it living in either place, really. Either way AbstractChains needs to be pulled.

[torfjelde]

Either way AbstractChains needs to be pulled.

Agreed. Do you know if there's anything stopping us from just depending on MCMCChains for the moment?

[cpfiffer]

Not that I'm aware of, but I've also been very out of the loop for two months.

[devmotion]

Shouldn't the issue that this PR tries to solve be fixed by #147 without having to deal with MCMCChains directly?

[torfjelde]

Ah, I did not see that! That should indeed fix the problem! But the "issue" of AbstractChains being a bit strange is still there 😕

But I guess we should just close this PR in favor of #147 then, and maybe make an issue over at the AbstractMCMC.jl regarding AbstractChains?

[yebai]

@torfjelde what's the status for this PR? If this is already fixed by #147, maybe consider closing it?

[torfjelde]

Closed in favour of #147 .