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(demo1(x), "sub1")
            b = @returned_quantities prefix(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)

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?

Fix #691

[codecov]

Codecov Report

Attention: Patch coverage is 93.75000% with 1 line in your changes missing coverage. Please review.

Project coverage is 77.78%. Comparing base (54691bf) to head (7aef65b).
Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
src/contexts.jl	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #696      +/-   ##
==========================================
- Coverage   79.22%   77.78%   -1.45%     
==========================================
  Files          30       30              
  Lines        4212     3938     -274     
==========================================
- Hits         3337     3063     -274     
  Misses        875      875              

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

[coveralls]

Pull Request Test Coverage Report for Build 11627641486

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

• 15 of 16 (93.75%) changed or added relevant lines in 4 files are covered.
• 62 unchanged lines in 2 files lost coverage.
• Overall coverage decreased (-2.2%) to 77.447%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/contexts.jl	4	5	80.0%

Files with Coverage Reduction	New Missed Lines	%
src/utils.jl	23	82.48%
src/varinfo.jl	39	80.24%

Totals
Change from base Build 11381380435:	-2.2%
Covered Lines:	3046
Relevant Lines:	3933

---

💛 - 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., #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

[mhauru]

For prefix/@prefix, maybe keep both but only export the macro? It sounds like unless you know what you are doing, you should use @prefix. And if you know what you're doing, you don't need it be exported. I do generally think it's a good idea to have a macro-free option available if possible.

For returned_quantities/@returned_quantities we still need both, because one is to be used outside of @model, the other inside, right? I forget what we concluded about this in our call, but I do worry users will mix the two up and get confusing errors.

[mhauru]

I found

@returned_quantities prefix=1 + 2 inner()

hard and unintuitive to parse. I think

@returned_quantities prefix=(1 + 2) inner()

would be much clearer. Not sure if this a documentation issue, or if we should disallow the former.

[torfjelde]

That's a documentation issue IMO, as this is not doing any special parsing but reliying on Julia's expression parsing.

[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.

[sunxd3]

I like the to_sampleable proposal and think name is something we can iterate on.
I would recommend creating an issue on documenting the transition from @submodel to to_sampleable before we merging this PR.

[mhauru]

After some discussion on Slack, I understand from @yebai that we are now going with this:

# To be implemented now
returned(model_instance, chain)  # Replaces generated_quantities(model_instance, chain)
x ~ to_sampleable(returned(submodel_instance))
x ~ to_submodel(submodel_instance)  # short-hand for the above

# To be implemented some time in the future
x ~ to_distribution(returned(submodel_instance))
x ~ to_sampleable(latent(submodel_instance))
x ~ to_distribution(latent(submodel_instance))

returned was introduced in the x ~ to_sampleable(returned(submodel_instance)) to make it more explicit that this is a distribution of the return values, not of the random variables used within the submodel. Outside of a @model, rand(returned(model)) would be equivalent to the current model(). latent would then do the complementary thing, giving access to the internal variables, and rand(latent(model)) would be equivalent to the current rand(model)

[torfjelde]

Alrighty! I've now moved the handling of the "submodels" into the tilde pipeline; specifically, I've added it to tilde_assume!! and tilde_observe!! which are specifically tilde functions which are not to be overridden by users (i.e. we treat this as a private-and-only-for-internal-overrides-usage).

This means that we:

1. Error when the user tries to use a submodel as a likelihood.
2. Support fix usage (which IMO makes complete sense).
3. Can technically handle cases where we want to sample the latents of the model instead of the return-values.

And the syntax is as follows:

@model outer() = a ~ returned(inner())

I've also added the option of doing

@model outer() = a ~ to_sampleable(returned(inner()))

as @yebai requested. However, this is unnecessary for returned(submodel), since in this case to_sampleable is just identity.

[mhauru]

Thanks for doing all the work for this @torfjelde. Could we also have a @test_throws for "~ with a model on the right-hand side of an observe statement is not supported"?

[torfjelde]

Could we also have a @test_throws for "~ with a model on the right-hand side of an observe statement is not supported"?

This is in the doctest:)

[torfjelde]

This can be generalized as we desire, e.g. if want to do something special with lantent(model), we can overload this to be true and then overload rand_like!!

[torfjelde]

Once we want "more" things to be allowed on right, we can easily deal with this by just generalizing is_rhs_model.

[torfjelde]

Accepted your suggestions @penelopeysm 👍

[penelopeysm]

Happy to approve once you're happy and if CI passes 😄

[penelopeysm]

Does this return true for Models themselves? It seems to me that it's only true for a ReturnedModelWrapper.

[torfjelde]

We dont't want to allow usage of Model on the RHS. So yeah, should probably rename the function. Long-term we probably just want something like is_valid_rhs_tilde(...) ,etc.

[torfjelde]

Okay, so the final call has been made and we'll introduce several wrappers:

• struct RerturnedModelWrapper: wraps a Model to indicate that it's a model over the return-values rather than its latents.
• abstract Distributional: a type indicating that something is distributional in some sense.
• struct Sampleable <: Distributional: a wrapper type that specifies that whatever it wraps is sampleable.

In combination with these types, we'll introduce the following functionality

to_sampleable(x) = Sampleable(x)
returned(model::Model) = ReturnedModelWrapper(model)
to_submodel(model::Model) = to_sampleable(returned(model))

At the moment, the Sampleable and Distributional (and their corresponding methods) doesn't have any particular purpose, but I believe @yebai has some plans for them in the future.

Moreover, the way we do it currently is that anything is allowed to on the RHS of a ~, as long as you override the correct methods.
But I think the plan (@yebai should confirm this) is that going forward we only allow subtypes of Distribution (from Distributions.jl) and Distributional (our type) to occur on the RHS of a ~ rather than specifying this by using dispatching / method overloading.

@yebai also mentioned that we want methods such as fix, condition, etc. to be supported for ReturnedModelWrapper, but I'll leave this for future PRs as a) we're not going do document that people can do this atm, so users will just use this inside a @model, and b) this, IMO, requires some more thought as to exactly how to execute. If we do it naively, we'll have to duplicate all of the necessary methods on Model, which is non-ideal. A likely approach is to introduce another abstract type AbstractModel which both Model and ReturnedModelWrapper inherits from; we could then define fix, etc. on this abstract type and only overload contextualize for the specific method in question.
But because this has dragged on for so long, let's leave this for another time.