Adds `@returned_quantities` macro

[torfjelde]

This adds the @returned_quantities macro as discussed @yebai @mhauru

This is meant to be a replacement for @submodel macro, but without the ability to do automatic prefixing. It ends up looking like

julia> @model function demo1(x)
           x ~ Normal()
           return 1 + abs(x)
       end;

julia> @model function demo2(x, y, z)
            a = @returned_quantities prefix="sub1" demo1(x)
            b = @returned_quantities prefix="sub2" demo1(y)
            return z ~ Uniform(-a, b)
       end;

julia> rand(demo2(missing, missing, 0.4))
(var"sub1.x" = 0.5865756059371534, var"sub2.x" = -0.25563799658500047)

Likely TODOs:

• [ ] Add deprecation warning to @submodel telling the user to use @returned_quantities.
• [ ] Do we do the renaming of generated_quantities to returned_quantities in this PR?

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 86.48%. Comparing base (82842bc) to head (b467c75). Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #696      +/-   ##
==========================================
+ Coverage   86.38%   86.48%   +0.09%     
==========================================
  Files          35       35              
  Lines        4180     4209      +29     
==========================================
+ Hits         3611     3640      +29     
  Misses        569      569              

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[coveralls]

Pull Request Test Coverage Report for Build 12144617184

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 34 of 36 (94.44%) changed or added relevant lines in 6 files are covered.
• 30 unchanged lines in 11 files lost coverage.
• Overall coverage increased (+0.09%) to 86.481%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/model.jl	14	16	87.5%
<!--	Total:	34	36

Files with Coverage Reduction	New Missed Lines	%
src/simple_varinfo.jl	1	86.08%
src/varnamedvector.jl	1	89.42%
src/sampler.jl	1	94.55%
src/debug_utils.jl	1	94.39%
ext/DynamicPPLEnzymeCoreExt.jl	1	0.0%
ext/DynamicPPLForwardDiffExt.jl	2	72.22%
src/utils.jl	3	83.16%
src/distribution_wrappers.jl	4	52.78%
src/contexts.jl	4	76.38%
src/values_as_in_model.jl	4	63.64%
<!--	Total:	30

Totals
Change from base Build 12138508432:	0.09%
Covered Lines:	3640
Relevant Lines:	4209

---

💛 - Coveralls

[yebai]

@torfjelde I suggest we change the prefix feature to a prefix_variables model operation (feel free to come up with better names). Then we could use the same functionality prefix_variables in more places, e.g.

# submodel prefixing
julia> @model function demo2(x, y, z)
            a = @returned_quantities prefix_variables(demo1(x), "sub1")
            b = @returned_quantities prefix_variables(demo1(y), "sub2")
            return z ~ Uniform(-a, b)
       end;

julia> rand(demo2(missing, missing, 0.4))
(var"sub1.x" = 0.5865756059371534, var"sub2.x" = -0.25563799658500047)

# rand prefixing 

julia> ret = rand(prefix_variables(demo1(1.), "prior_sample"))

# generated quantities / predict 

julia> returned_quantities(prefix_variables(demo1(1.), "generated_var_"), chain) 

This would also help unify the syntax of @generated_qunatities and generated_quantities- IIRC, the only difference between them is that generated_quantities lacks the prefixing/renaming feature.

This could be further unified with NamedDist in the future. See, e.g., https://github.com/TuringLang/DynamicPPL.jl/pull/414

[torfjelde]

We already have DynamicPPL.prefix, though this doesn't do exactly what you want here. We could easily just add

prefix(model::Model, x) = contextualize(model, PrefixContext(model.context, Symbol(x)))

or something as an additional definition.

However, I'm a bit worred about

1. It's quite verbose + a bit "too close to internals" for end-users.
2. To achieve the same performance guarantees that we have currently, we need to wrap everything in Val before calling prefix(model, ...) 😕 This seems non-ideal to me vs. the current approach.

[yebai]

It's quite verbose + a bit "too close to internals" for end-users.

I like the @returned_quantities(prefix(model, "prefix_")) syntax because it is

• less mysterious than @returned_quantities model "prefix_"
• all the other model operations could share this, e.g. rand(prefix(model, "prefix_")) to verify the effects of prefixing, which is very useful

prefix(model, x) is NOT any closer to internals than any other model operation APIs. They are the same, so this is not a problem.

To achieve the same performance guarantees that we have currently, we need to wrap everything in Val before calling prefix(model, ...) 😕 This seems non-ideal to me vs. the current approach.

Point taken, but this is very minor and a bit subjective.

[torfjelde]

Point taken, but this is very minor and a bit subjective.

But this means that the user needs to be careful and do prefix(model, Val{:whatever}()); if we just do prefix(model, :whatever), this will lead to type-instabilities. Do we really want to force end-users of Turing.jl to explicitly use Val? 😕

[yebai]

It is a standard Julia performance trick, so it is okay.

By default, we can print a performance warning message when users call prefix(model, x::String) or similiar.

[yebai]

I'm also happy to turn prefix into a macro: @prefix(model, :prefix_) if that helps. Then we could do

@returned_quantities @prefix(model, :prefix_)

[torfjelde]

Added a @prefix macro:) See the docstring of @returned_quantities for what it looks like 👍

[yebai]

Thanks, @torfjelde; I'm happy with the changes.

To minimise interface confusion (prefix vs. @prefix, and @returned_quantities vs. returned_quantities), shall we consider keeping only @prefix and @returned_quantities and depreciating generated_quantities and prefix?

Thoughts? @mhauru and @penelopeysm

[yebai]

For returned_quantities/@returned_quantities we still need both, because one is to be used outside of @model, the other inside, right?

generated_quantities allows users to fix model parameter values and/or accept MCMC chain objects. We can throw an error if users try to pass fixed parameter values or chain objects to @returned_quantities called within a model.

Then, @returned_quantities can match generated_quantities / returned_quantities exactly, thus allowing us to remove the generated_quantities / returned_quantities altogether.

[torfjelde]

Then, @returned_quantities can match generated_quantities / returned_quantities exactly, thus allowing us to remove the generated_quantities / returned_quantities altogether.

Just so we're all on the same page: @returned_quantities and returned_quantities will not match since the former only takes a single argument, while the other takes two, right? If so, then why would we want to raise explicit errors for incorrect arguments provided vs. just letting Julia raise the "not implemented error"?

[torfjelde]

Deprecated generated_quantities in favour of returned_quantities + removed the prefix=... argument for @prefix.

[yebai]

I think we could try to decompose returned_quantities(model, chain) to returned_quantitites(fix(model, chain)).

If that is viable, we need only keep the single-argument version of @returned_quantities and remove returned_quantities.

EDIT: Given their similar purposes, it is likely confusing for users to choose between @returned_quantities and returned_quantities.

[torfjelde]

I think we could try to decompose returned_quantities(model, chain) to returned_quantitites(fix(model, chain)).

This seems a bit weird to me? Here we would return multiple models from fix(model, chain)?

[yebai]

We still return one model: the generalised fix(model, chain) syntax of fix(model, value) is conceptually similar to "cut" distributions. We are replacing all model parameters' (prior) distributions with a discrete distribution of parameter values in chain.

[torfjelde]

We still return one model

But then what does the following piece of code do

rand(fix(model, chain))

?

[yebai]

rand(fix(model, chain))

That is, essentially, sampling uniformly a parameter realisation in the chain. If a parameter is not in the chain, we sample from its (conditional) prior as usual.

[yebai]

Also, before this PR, I had been thinking about fix(model, vector_of_values) to implement "cut" distributions, which seems very intuitive and flexible.

[torfjelde]

That is, essentially, sampling uniformly a parameter realisation in the chain. If a parameter is not in the chain, we sample from its (conditional) prior as usual.

This is different from what generated_quantities currently does though; it doesn't sample from chain but performs the conditional evaluation for every sample in chain.

Also, before this PR, I had been thinking about fix(model, vector_of_values) to implement "cut" distributions, which seems very intuitive and flexible.

But why not just do

fix(model, sample(chain))

?

[yebai]

This is different from what generated_quantities currently does though; it doesn't sample from chain but performs the conditional evaluation for every sample in chain.

I was referring to rand(fix(model, chain)). For @generated(fix(model, chain)), we iterate over all parameter realisations in the chain, instead of sampling them uniformly. Basically, fix(model, chain) simply returns a fixed_model, while its exact behaviour is controlled by the operation that acts on these fixed_model. So, rand and @returned_quantities could have different behaviour here, which is intuitive.

fix(model, sample(chain))

yes, for this simple case it is equivalent. But, think of fix(model, normalising_flow_approximation_of_theta), then we are actually replacing theta's prior with normalising_flow_approximation_of_theta. In short, fix could accept a single value, a vector of values, and a distribution object in its generalised form. Isn't that cool?

[torfjelde]

Isn't that cool?

Haha sure, but that also sounds very annoying to maintain 😅 I'm quite reluctant to have fix act in a bunch of different ways like this, when it's just replacing two lines of code 😕

[yebai]

fix(model, normalising_flow_approximation_of_theta)

Haha sure, but that also sounds very annoying to maintain 😅 I'm quite reluctant to have fix act in a bunch of different ways like this, when it's just replacing two lines of code 😕

This can be done in future PRs. It is not only cool but also very useful for compositional inference research.

@returned_quantities(fix(model, chain)), we iterate over all parameter realisations in the chain

For interface stability, let's consider pushing this (i.e. fix(model, chain)) through in this PR. This allows us to keep only the macro @returned_quantities and depreciate-then-remove generated_quantities/returned_quantities.

[torfjelde]

For interface stability, let's consider pushing this (i.e. fix(model, chain)) through in this PR.

Can you outline exactly what you want this to look like? IIUC you want @returned_quantities to work in two ways? Both as a macro inside a @model and outside @model (in which case it iterates over a chain)?

[torfjelde]

Might be useful to get some input from others on this too, e.g. @mhauru @penelopeysm @sunxd3 . I'm really skeptical of a single macro having multiple purposes like this. Overall I think it's just confusing for the user.

[yebai]

Can you outline exactly what you want this to look like? IIUC you want @returned_quantities to work in two ways? Both as a macro inside a @model and outside @model (in which case it iterates over a chain)?

yes, your understanding is correct.

I think this is quite intuitive because @returned_quantities's semantics for inside/outside model are exactly the same. In addition, replacing @returned_quantities(model, chain) with @returned_quantities(fix(model, chain)) also improves the compositionality of existing APIs.

Also happy to hear more thoughts from other team members.

[mhauru]

I'm in favour of only exporting @returned_quantities. When given a chain (outside of a @model), it would actually just call returned_quantities the function, but you would still call it as a macro. I think the average not-a-Julia-expert user would otherwise struggle to understand and remember when they are supposed to "put in the magic @", and might end up using the non-macro version within a model and getting confused by what the difference is.

I think I'm starting to get what @yebai means by the fix(model, chain) syntax. It would be a bit like saying fix(model, (x=>0.1)) except instead of a concrete value like 0.1, you could put in a random variable. In this case that random variable would be "x as sampled in chain". Correct me if I understood this wrong. Maybe you could also do things like fix(model, (x=>Normal()))?

However, I think having @returned_quantities(fix(model, chain)) return an iterable of the same length as chain is a bit confusing. It's an operation that doesn't depend on just the distribution defined by chain. It also depends on the particular way of representing that distribution as a chain. For instance, what would @returned_quantities(fix(model, normalising_flow_approximation_of_theta)) or @returned_quantities(fix(model, (x=>Normal()))) return? Presumably single values, rather than iterables?

It seems that the operation we want here is more like doing a map or broadcast of x -> @return_quantities(fix(model, x)) over chain. If it's not convenient to actually implement it with a map or broadcast syntax, then we could just call it @return_quantities(model, chain).

PS. After thought: Any love for @return_quantities.(model, chain)? It would be kinda nice if @returned_quantities would always return a single return value of model, like it does within a @model, and if you want an iterable, like with chain, then you would do some sort of map or broadcast. This would be in harmony with the general Julia convention of not overloading things like scalar + vector, and forcing people to use .+ instead.

[yebai]

Any love for @return_quantities.(model, chain)?

I like it: it is consistent with Julia's convention and very intuitive.

I think I'm starting to get what @yebai means by the fix(model, chain) syntax. It would be a bit like saying fix(model, (x=>0.1)) except instead of a concrete value like 0.1, you could put in a random variable.

This is correct.

In this case that random variable would be "x as sampled in chain". Correct me if I understood this wrong. Maybe you could also do things like fix(model, (x=>Normal()))?

There is something subtle here. Instead of "x as sampled in chain", fix(model, chain) means "x distributed as chain", where chain is a (multivariate) distribution with discrete support. Similarly fix(model, (x=>Normal())) means "x distributed as Normal". The exact computational meaning of this fix(model, chain) operation (e.g. sampling or computing logdensity) is determined by later operations acting on the fixed model. For example, we can perform HMC on a fixed model corresponding to nested MCMC for "cut" distributions. This kind of generalises the "do" operator to a "probabilistic do", which Rob and I are currently doing some research on. It is an exciting project.

[penelopeysm]

Personally, I would be wary of introducing too many overloads. There is a question of whether the benefits of doing this (convenience, elegance, etc - these are more immediate) outweigh the drawback of making the codebase harder to understand (which are often further down the road; it's naturally easy to think that we understand this PR because we're looking at it now, it's harder to say what happens if someone else reads it months or years later). I can't really make a judgment as to whether the changes in this PR are overall good or bad – I just want to make sure that this point has been considered :)

I really like the argument of compositionality, and I think that's a great aim. However, I don't see the need to use the same function name fix for it - we would still derive the same benefit if we used any other function name like fix_chain, and it would make code analysis down the line simpler perhaps?

[torfjelde]

Any love for @return_quantities.(model, chain)?

We can't broadcast macros, so not entirely certain what you mean by this?

We could ofc do returned_quantities.(model, chain), but that almost already works (we'd just need something like OrderedDictIterator(chain)).

I think I'm starting to get what @yebai means by the fix(model, chain) syntax. It would be a bit like saying fix(model, (x=>0.1)) except instead of a concrete value like 0.1, you could put in a random variable. In this case that random variable would be "x as sampled in chain". Correct me if I understood this wrong. Maybe you could also do things like fix(model, (x=>Normal()))?

But this is just a one-liner:

fix(model, sample(chain))