From 5c746c46d2b8b392ba51ac3b6b906c7a9fd77235 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 23 Oct 2024 09:52:39 -0600
Subject: [PATCH 01/63] Added `@returned_quantities` macro

---
 src/DynamicPPL.jl     |   1 +
 src/submodel_macro.jl | 182 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 183 insertions(+)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index a5d178125..dc27031fe 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -125,6 +125,7 @@ export AbstractVarInfo,
     # Convenience macros
     @addlogprob!,
     @submodel,
+    @returned_quantities,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 050bf31fc..b6c92078b 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -248,3 +248,185 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
         end
     end
 end
+
+"""
+    @returned_quantities [prefix=...] model
+
+Run `model` nested inside of another model and return the return-values of the `model`.
+
+Valid expressions for `prefix=...` are:
+- `prefix=false`: no prefix is used. This is the default.
+- `prefix=expression`: results in the prefix `Symbol(expression)`.
+
+Prefixing makes it possible to run the same model multiple times while keeping track of
+all random variables correctly, i.e. without name clashes.
+
+# Examples
+
+## Simple example
+```jldoctest submodel-returned-quantities; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y)
+            a = @returned_quantities(demo1(x))
+            return y ~ Uniform(0, a)
+       end;
+```
+
+When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
+```jldoctest submodel-return-quantities
+julia> vi = VarInfo(demo2(missing, 0.4));
+
+julia> @varname(x) in keys(vi)
+true
+```
+
+Variable `a` is not tracked since it can be computed from the random variable `x` that was
+tracked when running `demo1`:
+```jldoctest submodel-returned-quantities
+julia> @varname(a) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-return-quantities
+julia> x = vi[@varname(x)];
+
+julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
+true
+```
+
+## With prefixing
+```jldoctest submodel-return-quantities-prefix; setup=:(using Distributions)
+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;
+```
+
+When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
+`sub2.x` will be sampled:
+```jldoctest submodel-return-quantities-prefix
+julia> vi = VarInfo(demo2(missing, missing, 0.4));
+
+julia> @varname(var"sub1.x") in keys(vi)
+true
+
+julia> @varname(var"sub2.x") in keys(vi)
+true
+```
+
+Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
+`sub2.x` that were tracked when running `demo1`:
+```jldoctest submodel-return-quantities-prefix
+julia> @varname(a) in keys(vi)
+false
+
+julia> @varname(b) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-return-quantities-prefix
+julia> sub1_x = vi[@varname(var"sub1.x")];
+
+julia> sub2_x = vi[@varname(var"sub2.x")];
+
+julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
+
+julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
+
+julia> getlogp(vi) ≈ logprior + loglikelihood
+true
+```
+
+## Different ways of setting the prefix
+```jldoctest submodel-return-quantities-prefix-alts; setup=:(using DynamicPPL, Distributions)
+julia> @model inner() = x ~ Normal()
+inner (generic function with 2 methods)
+
+julia> # When `prefix` is unspecified, no prefix is used.
+       @model submodel_noprefix() = a = @returned_quantities inner()
+submodel_noprefix (generic function with 2 methods)
+
+julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
+true
+
+julia> # Explicitely don't use any prefix.
+       @model submodel_prefix_false() = a = @returned_quantities prefix=false inner()
+submodel_prefix_false (generic function with 2 methods)
+
+julia> @varname(x) in keys(VarInfo(submodel_prefix_false()))
+true
+
+julia> # Using a static string.
+       @model submodel_prefix_string() = a = @returned_quantities prefix="my prefix" inner()
+submodel_prefix_string (generic function with 2 methods)
+
+julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+true
+
+julia> # Using string interpolation.
+       @model submodel_prefix_interpolation() = a = @returned_quantities prefix="\$(nameof(inner()))" inner()
+submodel_prefix_interpolation (generic function with 2 methods)
+
+julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+true
+
+julia> # Or using some arbitrary expression.
+       @model submodel_prefix_expr() = a = @returned_quantities prefix=1 + 2 inner()
+submodel_prefix_expr (generic function with 2 methods)
+
+julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+true
+```
+"""
+macro returned_quantities(expr)
+    return returned_quantities_expr(:(prefix = false), expr)
+end
+
+macro returned_quantities(prefix_expr, expr)
+    return returned_quantities_expr(prefix_expr, expr)
+end
+
+"""
+    @returned_quantities_expr model
+
+Returns an expression that captures the return-values of a model in addition to the varinfo.
+"""
+function returned_quantities_expr(prefix_expr, expr, ctx=esc(:__context__))
+    prefix_left, prefix = getargs_assignment(prefix_expr)
+    if prefix_left !== :prefix
+        error("$(prefix_left) is not a valid kwarg")
+    end
+
+    # The user expects `@submodel ...` to return the
+    # return-value of the `...`, hence we need to capture
+    # the return-value and handle it correctly.
+    @gensym retval
+
+    # Prefix.
+    if prefix !== nothing
+        ctx = prefix_submodel_context(prefix, ctx)
+    end
+    return quote
+        # Evaluate the model and capture the return values + varinfo.
+        $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
+            $(esc(expr)), $(esc(:__varinfo__)), $(ctx)
+        )
+
+        # Return the return-value of the model.
+        $retval
+    end
+end

From 0b081b7dfecdcccd4c6eacd82cd868c588bb9b88 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 23 Oct 2024 10:10:29 -0600
Subject: [PATCH 02/63] Added `@returned_quantities` to the docs

---
 docs/src/api.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 638f6f3ee..1c7811c00 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,10 +14,11 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with [`@submodel`](@ref).
+One can nest models and call another model inside the model function with [`@submodel`](@ref) and [`@returned_quantities`](@ref).
 
 ```@docs
 @submodel
+@returned_quantities
 ```
 
 ### Type

From dc699a56482e7c74b5565853606dac0d91bada26 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 23 Oct 2024 11:15:42 -0600
Subject: [PATCH 03/63] Fixed names of doctests for `@returned_quantities`

---
 src/submodel_macro.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index b6c92078b..550a41d0e 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -277,7 +277,7 @@ julia> @model function demo2(x, y)
 ```
 
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
-```jldoctest submodel-return-quantities
+```jldoctest submodel-returned-quantities
 julia> vi = VarInfo(demo2(missing, 0.4));
 
 julia> @varname(x) in keys(vi)
@@ -293,7 +293,7 @@ false
 
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
-```jldoctest submodel-return-quantities
+```jldoctest submodel-returned-quantities
 julia> x = vi[@varname(x)];
 
 julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
@@ -301,7 +301,7 @@ true
 ```
 
 ## With prefixing
-```jldoctest submodel-return-quantities-prefix; setup=:(using Distributions)
+```jldoctest submodel-returned-quantities-prefix; setup=:(using Distributions)
 julia> @model function demo1(x)
            x ~ Normal()
            return 1 + abs(x)
@@ -316,7 +316,7 @@ julia> @model function demo2(x, y, z)
 
 When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
 `sub2.x` will be sampled:
-```jldoctest submodel-return-quantities-prefix
+```jldoctest submodel-returned-quantities-prefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
 
 julia> @varname(var"sub1.x") in keys(vi)
@@ -328,7 +328,7 @@ true
 
 Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
 `sub2.x` that were tracked when running `demo1`:
-```jldoctest submodel-return-quantities-prefix
+```jldoctest submodel-returned-quantities-prefix
 julia> @varname(a) in keys(vi)
 false
 
@@ -338,7 +338,7 @@ false
 
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
-```jldoctest submodel-return-quantities-prefix
+```jldoctest submodel-returned-quantities-prefix
 julia> sub1_x = vi[@varname(var"sub1.x")];
 
 julia> sub2_x = vi[@varname(var"sub2.x")];
@@ -352,7 +352,7 @@ true
 ```
 
 ## Different ways of setting the prefix
-```jldoctest submodel-return-quantities-prefix-alts; setup=:(using DynamicPPL, Distributions)
+```jldoctest submodel-returned-quantities-prefix-alts; setup=:(using DynamicPPL, Distributions)
 julia> @model inner() = x ~ Normal()
 inner (generic function with 2 methods)
 

From 7067695fedbfa9ed3419ac54015e655e6257c9a0 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 23 Oct 2024 21:39:14 -0600
Subject: [PATCH 04/63] Update src/submodel_macro.jl

Co-authored-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
---
 src/submodel_macro.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 550a41d0e..227097d7a 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -411,7 +411,7 @@ function returned_quantities_expr(prefix_expr, expr, ctx=esc(:__context__))
         error("$(prefix_left) is not a valid kwarg")
     end
 
-    # The user expects `@submodel ...` to return the
+    # The user expects `@returned_quantities ...` to return the
     # return-value of the `...`, hence we need to capture
     # the return-value and handle it correctly.
     @gensym retval

From 8cb07968d532f6e4514d35af8271ec80df4de4b9 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:06:51 -0400
Subject: [PATCH 05/63] Added `@prefix` macro which calls `prefix` with a `Val`
 argument to make things easier to basic users

---
 src/contexts.jl       | 28 ++++++++++++++++++++++++++++
 src/submodel_macro.jl | 17 ++++++++---------
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index 53b454df6..8df3e1ab8 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -294,6 +294,34 @@ function prefix(::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
     end
 end
 
+"""
+    prefix(model::Model, x)
+
+Return `model` but with all random variables prefixed by `x`.
+
+If `x` is known at compile-time, use `Val{x}()` to avoid runtime overheads for prefixing.
+"""
+prefix(model::Model, x) = contextualize(model, PrefixContext{Symbol(x)}(model.context))
+function prefix(model::Model, ::Val{x}) where {x}
+    contextualize(model, PrefixContext{Symbol(x)}(model.context))
+end
+
+"""
+    @prefix(model, prefix_expr)
+
+Return `model` but with all random variables prefixed by `prefix_expr`.
+
+The result of `prefix_expr` must will be converted to a `Symbol` and used as the prefix.
+
+!!! note
+    This is effectively just a convenience macro for the method [`prefix(::Model, x)`](@ref),
+    which automatically converts the result of `prefix_expr` into a `Val` to avoid runtime overheads
+    for static prefixes. For more control over the prefixing, use the method directly.
+"""
+macro prefix(model, prefix_expr)
+    return :($prefix($(esc(model)), $Val{$(esc(prefix_expr))}()))
+end
+
 struct ConditionContext{Values,Ctx<:AbstractContext} <: AbstractContext
     values::Values
     context::Ctx
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 550a41d0e..0de49a6e8 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -250,16 +250,15 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
 end
 
 """
-    @returned_quantities [prefix=...] model
+    @returned_quantities model
 
 Run `model` nested inside of another model and return the return-values of the `model`.
 
-Valid expressions for `prefix=...` are:
-- `prefix=false`: no prefix is used. This is the default.
-- `prefix=expression`: results in the prefix `Symbol(expression)`.
-
-Prefixing makes it possible to run the same model multiple times while keeping track of
-all random variables correctly, i.e. without name clashes.
+!!! warning
+    It's generally recommended to use [`prefix(::Model, x)`](@ref) or
+    [`@prefix(model, prefix_expr)`](@ref) in combination with `@returned_quantities`
+    to ensure that the variables in `model` are unique and do not clash with other variables in the
+    parent model or in other submodels.
 
 # Examples
 
@@ -308,8 +307,8 @@ julia> @model function demo1(x)
        end;
 
 julia> @model function demo2(x, y, z)
-            a = @returned_quantities prefix="sub1" demo1(x)
-            b = @returned_quantities prefix="sub2" demo1(y)
+            a = @returned_quantities prefix(demo1(x), Val{:sub1}())
+            b = @returned_quantities prefix(demo1(y), Val{:sub2}())
             return z ~ Uniform(-a, b)
        end;
 ```

From 2d887c970a995ed8a31d5161c16a511286f95ca6 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:10:35 -0400
Subject: [PATCH 06/63] Convert the result of `prefix_expr` in `@prefix` into a
 `Sybmol` before wrapping in `Val`

---
 src/contexts.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index 8df3e1ab8..104965961 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -319,7 +319,7 @@ The result of `prefix_expr` must will be converted to a `Symbol` and used as the
     for static prefixes. For more control over the prefixing, use the method directly.
 """
 macro prefix(model, prefix_expr)
-    return :($prefix($(esc(model)), $Val{$(esc(prefix_expr))}()))
+    return :($prefix($(esc(model)), $Val{$Symbol($(esc(prefix_expr)))}()))
 end
 
 struct ConditionContext{Values,Ctx<:AbstractContext} <: AbstractContext

From 692cfffd3794fa98dfbc0f09d9daaf52cf178d67 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:11:10 -0400
Subject: [PATCH 07/63] Export `prefix` and `@prefix`

---
 src/DynamicPPL.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index dc27031fe..d2643cb1d 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -126,6 +126,8 @@ export AbstractVarInfo,
     @addlogprob!,
     @submodel,
     @returned_quantities,
+    prefix,
+    @prefix,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,

From 32fd6b964e4b18dc93161f7aa35ffcab4920f3e1 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:11:25 -0400
Subject: [PATCH 08/63] Updated docstring for `@returned_quantities`

---
 src/submodel_macro.jl | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 0de49a6e8..71410bc47 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -307,8 +307,8 @@ julia> @model function demo1(x)
        end;
 
 julia> @model function demo2(x, y, z)
-            a = @returned_quantities prefix(demo1(x), Val{:sub1}())
-            b = @returned_quantities prefix(demo1(y), Val{:sub2}())
+            a = @returned_quantities @prefix(demo1(x), :sub1)
+            b = @returned_quantities @prefix(demo1(y), :sub2)
             return z ~ Uniform(-a, b)
        end;
 ```
@@ -362,22 +362,15 @@ submodel_noprefix (generic function with 2 methods)
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
 true
 
-julia> # Explicitely don't use any prefix.
-       @model submodel_prefix_false() = a = @returned_quantities prefix=false inner()
-submodel_prefix_false (generic function with 2 methods)
-
-julia> @varname(x) in keys(VarInfo(submodel_prefix_false()))
-true
-
 julia> # Using a static string.
-       @model submodel_prefix_string() = a = @returned_quantities prefix="my prefix" inner()
+       @model submodel_prefix_string() = a = @returned_quantities @prefix inner() "my prefix"
 submodel_prefix_string (generic function with 2 methods)
 
 julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
 true
 
 julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a = @returned_quantities prefix="\$(nameof(inner()))" inner()
+       @model submodel_prefix_interpolation() = a = @returned_quantities @prefix inner() "\$(nameof(inner()))"
 submodel_prefix_interpolation (generic function with 2 methods)
 
 julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))

From 5478fb346674e0bd46aa2983c1d515c4e308c7ef Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:11:39 -0400
Subject: [PATCH 09/63] Fixed bug in `rand` for `Model` where it would
 duplicate the non-leaf contexts in `model.context`

---
 src/model.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index 2a1a6db88..580a9f9c7 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1051,7 +1051,9 @@ function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
         evaluate!!(
             model,
             SimpleVarInfo{Float64}(OrderedDict()),
-            SamplingContext(rng, SampleFromPrior(), model.context),
+            # NOTE: Use `leafcontext` here so we a) avoid overriding the leaf context of `model`,
+            # and b) avoid double-stacking the parent contexts.
+            SamplingContext(rng, SampleFromPrior(), leafcontext(model.context)),
         ),
     )
     return values_as(x, T)

From 9e0730fce4f1f6cc60355e84a058e6f0810426fc Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:45:16 -0400
Subject: [PATCH 10/63] Update src/contexts.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 src/contexts.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index 104965961..22c553442 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -303,7 +303,7 @@ If `x` is known at compile-time, use `Val{x}()` to avoid runtime overheads for p
 """
 prefix(model::Model, x) = contextualize(model, PrefixContext{Symbol(x)}(model.context))
 function prefix(model::Model, ::Val{x}) where {x}
-    contextualize(model, PrefixContext{Symbol(x)}(model.context))
+    return contextualize(model, PrefixContext{Symbol(x)}(model.context))
 end
 
 """

From cc3af4688d150385e399f3303eb3ee6c3a82d881 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 29 Oct 2024 18:55:39 -0400
Subject: [PATCH 11/63] Added `prefix` and `@prefix` to docs

---
 docs/src/api.md | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/docs/src/api.md b/docs/src/api.md
index 1c7811c00..fcd38ded2 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -21,6 +21,13 @@ One can nest models and call another model inside the model function with [`@sub
 @returned_quantities
 ```
 
+In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
+
+```@docs
+prefix(::DynamicPPL.Model, x)
+@prefix
+```
+
 ### Type
 
 A [`Model`](@ref) can be created by calling the model function, as defined by [`@model`](@ref).

From 720053a77bacf35cbe6a8a71a570688efa9c4aaa Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 11:37:16 +0100
Subject: [PATCH 12/63] removed the prefix=... syntax for
 `@returned_quantities`

---
 src/submodel_macro.jl | 30 ++++++++++--------------------
 1 file changed, 10 insertions(+), 20 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 6fdbb637a..ff93b4020 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -363,21 +363,21 @@ julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
 true
 
 julia> # Using a static string.
-       @model submodel_prefix_string() = a = @returned_quantities @prefix inner() "my prefix"
+       @model submodel_prefix_string() = a = @returned_quantities @prefix(inner(), "my prefix")
 submodel_prefix_string (generic function with 2 methods)
 
 julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
 true
 
 julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a = @returned_quantities @prefix inner() "\$(nameof(inner()))"
+       @model submodel_prefix_interpolation() = a = @returned_quantities @prefix(inner(), "\$(nameof(inner()))")
 submodel_prefix_interpolation (generic function with 2 methods)
 
 julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
 true
 
 julia> # Or using some arbitrary expression.
-       @model submodel_prefix_expr() = a = @returned_quantities prefix=1 + 2 inner()
+       @model submodel_prefix_expr() = a = @returned_quantities @prefix(inner(), 1 + 2)
 submodel_prefix_expr (generic function with 2 methods)
 
 julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
@@ -385,37 +385,27 @@ true
 ```
 """
 macro returned_quantities(expr)
-    return returned_quantities_expr(:(prefix = false), expr)
-end
-
-macro returned_quantities(prefix_expr, expr)
-    return returned_quantities_expr(prefix_expr, expr)
+    return returned_quantities_expr(expr)
 end
 
 """
-    @returned_quantities_expr model
+    returned_quantities_expr(model_expr[, ctx_expr])
 
 Returns an expression that captures the return-values of a model in addition to the varinfo.
-"""
-function returned_quantities_expr(prefix_expr, expr, ctx=esc(:__context__))
-    prefix_left, prefix = getargs_assignment(prefix_expr)
-    if prefix_left !== :prefix
-        error("$(prefix_left) is not a valid kwarg")
-    end
 
+!!! warning
+    This is only meant to be used in the body of `@model`.
+"""
+function returned_quantities_expr(model_expr, ctx_expr=esc(:__context__))
     # The user expects `@returned_quantities ...` to return the
     # return-value of the `...`, hence we need to capture
     # the return-value and handle it correctly.
     @gensym retval
 
-    # Prefix.
-    if prefix !== nothing
-        ctx = prefix_submodel_context(prefix, ctx)
-    end
     return quote
         # Evaluate the model and capture the return values + varinfo.
         $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
-            $(esc(expr)), $(esc(:__varinfo__)), $(ctx)
+            $(esc(model_expr)), $(esc(:__varinfo__)), $(ctx_expr)
         )
 
         # Return the return-value of the model.

From fe0403f52fde8bfaa4811113766036955702fa39 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 11:40:42 +0100
Subject: [PATCH 13/63] added deprecation.jl + deprecated
 `generated_quantities` in favour of `returned_quantities`

---
 src/DynamicPPL.jl |  2 ++
 src/deprecated.jl |  1 +
 src/model.jl      | 16 ++++++++--------
 3 files changed, 11 insertions(+), 8 deletions(-)
 create mode 100644 src/deprecated.jl

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index d2643cb1d..4dc8dd06b 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -199,6 +199,8 @@ include("values_as_in_model.jl")
 include("debug_utils.jl")
 using .DebugUtils
 
+include("deprecated.jl")
+
 if !isdefined(Base, :get_extension)
     using Requires
 end
diff --git a/src/deprecated.jl b/src/deprecated.jl
new file mode 100644
index 000000000..32492a9d9
--- /dev/null
+++ b/src/deprecated.jl
@@ -0,0 +1 @@
+@deprecate generated_quantities returned_quantities
diff --git a/src/model.jl b/src/model.jl
index 580a9f9c7..2372e3389 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1206,9 +1206,9 @@ function Distributions.loglikelihood(model::Model, chain::AbstractMCMC.AbstractC
 end
 
 """
-    generated_quantities(model::Model, parameters::NamedTuple)
-    generated_quantities(model::Model, values, keys)
-    generated_quantities(model::Model, values, keys)
+    returned_quantities(model::Model, parameters::NamedTuple)
+    returned_quantities(model::Model, values, keys)
+    returned_quantities(model::Model, values, keys)
 
 Execute `model` with variables `keys` set to `values` and return the values returned by the `model`.
 
@@ -1233,18 +1233,18 @@ julia> model = demo(randn(10));
 
 julia> parameters = (; s = 1.0, m_shifted=10.0);
 
-julia> generated_quantities(model, parameters)
+julia> returned_quantities(model, parameters)
 (0.0,)
 
-julia> generated_quantities(model, values(parameters), keys(parameters))
+julia> returned_quantities(model, values(parameters), keys(parameters))
 (0.0,)
 ```
 """
-function generated_quantities(model::Model, parameters::NamedTuple)
+function returned_quantities(model::Model, parameters::NamedTuple)
     fixed_model = fix(model, parameters)
     return fixed_model()
 end
 
-function generated_quantities(model::Model, values, keys)
-    return generated_quantities(model, NamedTuple{keys}(values))
+function returned_quantities(model::Model, values, keys)
+    return returned_quantities(model, NamedTuple{keys}(values))
 end

From 55b95a116a803b46d0d8c67038e389aa2ad48f1b Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 11:42:15 +0100
Subject: [PATCH 14/63] removed export of `prefix` and `generated_quantities`
 (the latter is exported by the deprecation macro)

---
 src/DynamicPPL.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 4dc8dd06b..0747a4374 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -86,7 +86,7 @@ export AbstractVarInfo,
     Model,
     getmissings,
     getargnames,
-    generated_quantities,
+    returned_quantities,
     extract_priors,
     values_as_in_model,
     # Samplers
@@ -126,7 +126,6 @@ export AbstractVarInfo,
     @addlogprob!,
     @submodel,
     @returned_quantities,
-    prefix,
     @prefix,
     value_iterator_from_chain,
     check_model,

From 34fb6bd491fc2278eae3c2e2d12f433b12dfeb46 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 12:07:53 +0100
Subject: [PATCH 15/63] updated `DynamicPPLMCMCChainsExt` to define
 `returned_quantities`

---
 ext/DynamicPPLMCMCChainsExt.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
index c91fb1fe0..1b2fa0059 100644
--- a/ext/DynamicPPLMCMCChainsExt.jl
+++ b/ext/DynamicPPLMCMCChainsExt.jl
@@ -43,7 +43,7 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
 end
 
 """
-    generated_quantities(model::Model, chain::MCMCChains.Chains)
+    returned_quantities(model::Model, chain::MCMCChains.Chains)
 
 Execute `model` for each of the samples in `chain` and return an array of the values
 returned by the `model` for each sample.
@@ -63,7 +63,7 @@ m = demo(data)
 chain = sample(m, alg, n)
 # To inspect the `interesting_quantity(θ, x)` where `θ` is replaced by samples
 # from the posterior/`chain`:
-generated_quantities(m, chain) # <= results in a `Vector` of returned values
+returned_quantities(m, chain) # <= results in a `Vector` of returned values
                                #    from `interesting_quantity(θ, x)`
 ```
 ## Concrete (and simple)
@@ -87,7 +87,7 @@ julia> model = demo(randn(10));
 
 julia> chain = sample(model, MH(), 10);
 
-julia> generated_quantities(model, chain)
+julia> returned_quantities(model, chain)
 10×1 Array{Tuple{Float64},2}:
  (2.1964758025119338,)
  (2.1964758025119338,)
@@ -101,7 +101,7 @@ julia> generated_quantities(model, chain)
  (-0.16489786710222099,)
 ```
 """
-function DynamicPPL.generated_quantities(
+function DynamicPPL.returned_quantities(
     model::DynamicPPL.Model, chain_full::MCMCChains.Chains
 )
     chain = MCMCChains.get_sections(chain_full, :parameters)

From 9a7e18fd9d43a1bf23390340706ecd2f3edb23d7 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 12:08:42 +0100
Subject: [PATCH 16/63] updated docs

---
 docs/src/api.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index fcd38ded2..3e4313aa0 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -126,10 +126,10 @@ It is possible to manually increase (or decrease) the accumulated log density fr
 @addlogprob!
 ```
 
-Return values of the model function for a collection of samples can be obtained with [`generated_quantities`](@ref).
+Return values of the model function for a collection of samples can be obtained with [`returned_quantities`](@ref).
 
 ```@docs
-generated_quantities
+returned_quantities
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using

From 7aef65b700b3bd1fd4f676979f7ddb569c67982b Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 1 Nov 2024 11:00:46 +0100
Subject: [PATCH 17/63] Update docs/src/api.md

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>
---
 docs/src/api.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 3e4313aa0..4b04a0fbf 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -24,7 +24,6 @@ One can nest models and call another model inside the model function with [`@sub
 In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
 
 ```@docs
-prefix(::DynamicPPL.Model, x)
 @prefix
 ```
 

From 5ee727b0264b503dfaf50bc1f05efc77133771a1 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:14:00 +0100
Subject: [PATCH 18/63] improved docstring for `prefix` and `@prefix`

---
 src/contexts.jl | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/contexts.jl b/src/contexts.jl
index 22c553442..171246610 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -300,6 +300,22 @@ end
 Return `model` but with all random variables prefixed by `x`.
 
 If `x` is known at compile-time, use `Val{x}()` to avoid runtime overheads for prefixing.
+
+# Examples
+
+```jldoctest
+julia> using DynamicPPL: prefix
+
+julia> @model demo() = x ~ Dirac(1)
+demo (generic function with 2 methods)
+
+julia> rand(prefix(demo(), :my_prefix))
+(var"my_prefix.x" = 1,)
+
+julia> # One can also use `Val` to avoid runtime overheads.
+       rand(prefix(demo(), Val(:my_prefix)))
+(var"my_prefix.x" = 1,)
+```
 """
 prefix(model::Model, x) = contextualize(model, PrefixContext{Symbol(x)}(model.context))
 function prefix(model::Model, ::Val{x}) where {x}
@@ -314,9 +330,19 @@ Return `model` but with all random variables prefixed by `prefix_expr`.
 The result of `prefix_expr` must will be converted to a `Symbol` and used as the prefix.
 
 !!! note
-    This is effectively just a convenience macro for the method [`prefix(::Model, x)`](@ref),
+    This is effectively just a convenience macro for the method [`DynamicPPL.prefix(::Model, x)`](@ref),
     which automatically converts the result of `prefix_expr` into a `Val` to avoid runtime overheads
     for static prefixes. For more control over the prefixing, use the method directly.
+
+# Examples
+
+```jldoctest
+julia> @model demo() = x ~ Dirac(1)
+demo (generic function with 2 methods)
+
+julia> rand(@prefix(demo(), :my_prefix))
+(var"my_prefix.x" = 1,)
+```
 """
 macro prefix(model, prefix_expr)
     return :($prefix($(esc(model)), $Val{$Symbol($(esc(prefix_expr)))}()))

From d92141c7b1a6a249c2f5df19bc2a9c44ad8ef94b Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:20:42 +0100
Subject: [PATCH 19/63] added `@returned_quantities` macro taking two arguments
 + removed `returned_quantities` from exported functions

---
 src/DynamicPPL.jl |  1 -
 src/model.jl      | 12 ++++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 0747a4374..2bb7556ba 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -86,7 +86,6 @@ export AbstractVarInfo,
     Model,
     getmissings,
     getargnames,
-    returned_quantities,
     extract_priors,
     values_as_in_model,
     # Samplers
diff --git a/src/model.jl b/src/model.jl
index 2372e3389..442cf6cfa 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1248,3 +1248,15 @@ end
 function returned_quantities(model::Model, values, keys)
     return returned_quantities(model, NamedTuple{keys}(values))
 end
+
+"""
+    @returned_quantities(model, input)
+
+Execute `model` and extract the return-values of `model` for `input`.
+
+!!! note
+    This macro is in fact a simple wrapper around the method [`DynamicPPL.returned_quantities`](@ref).
+"""
+macro returned_quantities(model_expr, input_expr)
+    return :($returned_quantities($(esc(model_expr)), $(esc(input_expr))))
+end

From 64b519d7f82f318fe1252e41c277e270c2b76777 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:21:57 +0100
Subject: [PATCH 20/63] updated docs to reflect the new two-argument
 `@returned_quantities`

---
 docs/src/api.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 4b04a0fbf..f6da8848a 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,11 +14,11 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with [`@submodel`](@ref) and [`@returned_quantities`](@ref).
+One can nest models and call another model inside the model function with [`@submodel`](@ref) and [`@returned_quantities(model)`](@ref).
 
 ```@docs
 @submodel
-@returned_quantities
+@returned_quantities(model)
 ```
 
 In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
@@ -125,10 +125,10 @@ It is possible to manually increase (or decrease) the accumulated log density fr
 @addlogprob!
 ```
 
-Return values of the model function for a collection of samples can be obtained with [`returned_quantities`](@ref).
+Return values of the model function for a collection of samples can be obtained with [`@returned_quantities`](@ref).
 
 ```@docs
-returned_quantities
+@returned_quantities(model, input)
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using

From 1b48f6567e0263d071eae177d8048a190849f1ae Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:37:29 +0100
Subject: [PATCH 21/63] added depwarn to `@submodel` macro

---
 src/submodel_macro.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index ff93b4020..e994928af 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -207,6 +207,8 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` is deprecated, use `@returned_quantities model` instead."
+
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
     if prefix_left !== :prefix
@@ -225,6 +227,9 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
     return if args_assign === nothing
         ctx = prefix_submodel_context(prefix, ctx)
         quote
+            # Raise deprecation warning to let user know that we recommend using `@returned_quantities`.
+            $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
+
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
                 $(esc(expr)), $(esc(:__varinfo__)), $(ctx)
             )
@@ -241,6 +246,9 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
             )
         end
         quote
+            # Raise deprecation warning to let user know that we recommend using `@returned_quantities`.
+            $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
+
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
                 $(esc(R)), $(esc(:__varinfo__)), $(ctx)
             )

From db2102cd3307d6ff6636f2211069cb178021f7b7 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:38:26 +0100
Subject: [PATCH 22/63] fixed reference

---
 src/submodel_macro.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index e994928af..a831b1c11 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -263,7 +263,7 @@ end
 Run `model` nested inside of another model and return the return-values of the `model`.
 
 !!! warning
-    It's generally recommended to use [`prefix(::Model, x)`](@ref) or
+    It's generally recommended to use [`DynamicPPL.prefix(::Model, x)`](@ref) or
     [`@prefix(model, prefix_expr)`](@ref) in combination with `@returned_quantities`
     to ensure that the variables in `model` are unique and do not clash with other variables in the
     parent model or in other submodels.

From da95aba44eb89ca2613fee55b1bdba56d504540b Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 16:38:46 +0100
Subject: [PATCH 23/63] fixed reference to `@prefix` in `@returned_quantities`
 macro

---
 src/submodel_macro.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index a831b1c11..3ee9eaa26 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -263,7 +263,7 @@ end
 Run `model` nested inside of another model and return the return-values of the `model`.
 
 !!! warning
-    It's generally recommended to use [`DynamicPPL.prefix(::Model, x)`](@ref) or
+    It's generally recommended to use [`@prefix(::Model, input)`](@ref) or
     [`@prefix(model, prefix_expr)`](@ref) in combination with `@returned_quantities`
     to ensure that the variables in `model` are unique and do not clash with other variables in the
     parent model or in other submodels.

From c8d567fcbd30332aa8a22684db5927d7738e5b0c Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 6 Nov 2024 17:29:41 +0100
Subject: [PATCH 24/63] actually fixed doc references

---
 docs/src/api.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/src/api.md b/docs/src/api.md
index f6da8848a..94688b481 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -25,6 +25,7 @@ In the context of nesting models, it's also useful to prefix the variables in su
 
 ```@docs
 @prefix
+DynamicPPL.prefix
 ```
 
 ### Type
@@ -129,6 +130,7 @@ Return values of the model function for a collection of samples can be obtained
 
 ```@docs
 @returned_quantities(model, input)
+DynamicPPL.returned_quantities
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using

From d477137e79c57cd355dc7c2bace1ee5b47e88ff6 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 8 Nov 2024 11:44:02 +0100
Subject: [PATCH 25/63] updated doctests for `@submodel` to include the depwarn
 + added warning regarding deprecation of `@submodel`

---
 ext/DynamicPPLMCMCChainsExt.jl |  4 ++--
 src/model.jl                   |  2 ++
 src/submodel_macro.jl          | 17 +++++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
index 1b2fa0059..eb9745552 100644
--- a/ext/DynamicPPLMCMCChainsExt.jl
+++ b/ext/DynamicPPLMCMCChainsExt.jl
@@ -68,7 +68,7 @@ returned_quantities(m, chain) # <= results in a `Vector` of returned values
 ```
 ## Concrete (and simple)
 ```julia
-julia> using DynamicPPL, Turing
+julia> using Turing
 
 julia> @model function demo(xs)
            s ~ InverseGamma(2, 3)
@@ -87,7 +87,7 @@ julia> model = demo(randn(10));
 
 julia> chain = sample(model, MH(), 10);
 
-julia> returned_quantities(model, chain)
+julia> DynamicPPL.returned_quantities(model, chain)
 10×1 Array{Tuple{Float64},2}:
  (2.1964758025119338,)
  (2.1964758025119338,)
diff --git a/src/model.jl b/src/model.jl
index 442cf6cfa..39497d947 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1218,6 +1218,8 @@ If a `NamedTuple` is given, `keys=keys(parameters)` and `values=values(parameter
 ```jldoctest
 julia> using DynamicPPL, Distributions
 
+julia> using DynamicPPL: returned_quantities
+
 julia> @model function demo(xs)
            s ~ InverseGamma(2, 3)
            m_shifted ~ Normal(10, √s)
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 3ee9eaa26..98dbf6201 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -4,6 +4,10 @@
 
 Run a Turing `model` nested inside of a Turing model.
 
+!!! warning
+    This is deprecated and will be removed in a future release.
+    Use [`@returned_quantities(model)`](@ref) instead.
+
 # Examples
 
 ```jldoctest submodel; setup=:(using Distributions)
@@ -21,6 +25,9 @@ julia> @model function demo2(x, y)
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
 ```jldoctest submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
+┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+│   caller = ip:0x0
+└ @ Core :-1
 
 julia> @varname(x) in keys(vi)
 true
@@ -62,6 +69,10 @@ Valid expressions for `prefix=...` are:
 The prefix makes it possible to run the same Turing model multiple times while
 keeping track of all random variables correctly.
 
+!!! warning
+    This is deprecated and will be removed in a future release.
+    Use [`@returned_quantities`](@ref) combined with [`@prefix`](@ref) instead.
+
 # Examples
 ## Example models
 ```jldoctest submodelprefix; setup=:(using Distributions)
@@ -81,6 +92,9 @@ When we sample from the model `demo2(missing, missing, 0.4)` random variables `s
 `sub2.x` will be sampled:
 ```jldoctest submodelprefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
+┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+│   caller = ip:0x0
+└ @ Core :-1
 
 julia> @varname(var"sub1.x") in keys(vi)
 true
@@ -124,6 +138,9 @@ julia> # When `prefix` is unspecified, no prefix is used.
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
+┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # Explicitely don't use any prefix.

From bf35de4506071a3fa102c77afbc8f1f53db6ec61 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 15 Nov 2024 20:28:46 +0100
Subject: [PATCH 26/63] added `to_sampleable` and limited `~` handling for
 submodels

---
 src/DynamicPPL.jl |  1 +
 src/compiler.jl   | 52 +++++++++++++++++++++++++++++------------------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 2bb7556ba..950c1a8f5 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -129,6 +129,7 @@ export AbstractVarInfo,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,
+    to_sampleable,
     # Deprecated.
     @logprob_str,
     @prob_str
diff --git a/src/compiler.jl b/src/compiler.jl
index 90220cbf5..404766bc5 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -1,5 +1,11 @@
 const INTERNALNAMES = (:__model__, :__context__, :__varinfo__)
 
+struct SampleableModelWrapper{M}
+    model::M
+end
+
+to_sampleable(model::DynamicPPL.Model) = SampleableModelWrapper(model)
+
 """
     need_concretize(expr)
 
@@ -427,28 +433,34 @@ function generate_tilde(left, right)
     # more selective with our escape. Until that's the case, we remove them all.
     return quote
         $dist = $right
-        $vn = $(DynamicPPL.resolve_varnames)(
-            $(AbstractPPL.drop_escape(varname(left, need_concretize(left)))), $dist
-        )
-        $isassumption = $(DynamicPPL.isassumption(left, vn))
-        if $(DynamicPPL.isfixed(left, vn))
-            $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
-        elseif $isassumption
-            $(generate_tilde_assume(left, dist, vn))
-        else
-            # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
-            if !$(DynamicPPL.inargnames)($vn, __model__)
-                $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
-            end
 
-            $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
-                __context__,
-                $(DynamicPPL.check_tilde_rhs)($dist),
-                $(maybe_view(left)),
-                $vn,
-                __varinfo__,
+        if $dist isa $(SampleableModelWrapper)
+            $left, __varinfo__ = $(_evaluate!!)($dist.model, __varinfo__, __context__)
+            $left
+        else
+            $vn = $(DynamicPPL.resolve_varnames)(
+                $(AbstractPPL.drop_escape(varname(left, need_concretize(left)))), $dist
             )
-            $value
+            $isassumption = $(DynamicPPL.isassumption(left, vn))
+            if $(DynamicPPL.isfixed(left, vn))
+                $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
+            elseif $isassumption
+                $(generate_tilde_assume(left, dist, vn))
+            else
+                # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
+                if !$(DynamicPPL.inargnames)($vn, __model__)
+                    $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
+                end
+
+                $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
+                    __context__,
+                    $(DynamicPPL.check_tilde_rhs)($dist),
+                    $(maybe_view(left)),
+                    $vn,
+                    __varinfo__,
+                )
+                $value
+            end
         end
     end
 end

From 0f2062465156e959c22a60caae4e3fdaaa7c2a23 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 15 Nov 2024 21:01:46 +0100
Subject: [PATCH 27/63] added docs to `to_sampleable` + removed the unnecessary
 macro exports that we no longer need

---
 docs/src/api.md       |  17 +++--
 src/DynamicPPL.jl     |   4 +-
 src/compiler.jl       | 143 +++++++++++++++++++++++++++++++++-
 src/contexts.jl       |  26 -------
 src/model.jl          |  12 ---
 src/submodel_macro.jl | 174 ++----------------------------------------
 6 files changed, 159 insertions(+), 217 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index c7913278b..962939c2c 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,18 +14,22 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with [`@submodel`](@ref) and [`@returned_quantities(model)`](@ref).
+One can nest models and call another model inside the model function with `left ~ to_sampleable(model)`.
+
+```@docs
+to_sampleable
+```
+
+In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`to_sampleable(model)`](@ref)
 
 ```@docs
 @submodel
-@returned_quantities(model)
 ```
 
 In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
 
 ```@docs
-@prefix
-DynamicPPL.prefix
+prefix
 ```
 
 ### Type
@@ -126,11 +130,10 @@ It is possible to manually increase (or decrease) the accumulated log density fr
 @addlogprob!
 ```
 
-Return values of the model function for a collection of samples can be obtained with [`@returned_quantities`](@ref).
+Return values of the model function for a collection of samples can be obtained with [`returned_quantities`](@ref).
 
 ```@docs
-@returned_quantities(model, input)
-DynamicPPL.returned_quantities
+returned_quantities
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 950c1a8f5..0334bd466 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -121,11 +121,11 @@ export AbstractVarInfo,
     decondition,
     fix,
     unfix,
+    prefix,
+    returned_quantities,
     # Convenience macros
     @addlogprob!,
     @submodel,
-    @returned_quantities,
-    @prefix,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,
diff --git a/src/compiler.jl b/src/compiler.jl
index 404766bc5..f19cd5c75 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -4,7 +4,148 @@ struct SampleableModelWrapper{M}
     model::M
 end
 
-to_sampleable(model::DynamicPPL.Model) = SampleableModelWrapper(model)
+"""
+    to_sampleable(model::Model)
+
+Return a wrapper around `model` which indicates that this model can only be sampled from.
+
+This is mainly meant to be used on the right-hand side of a `~` operator to indicate that
+the model can be sampled from but not necessarily evaluated for its log density.
+
+!!! warning
+    Note that other operations that one typically associate with expressions of the form `left ~ right`
+    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `to_sampleable`.
+
+!!! warning
+    It's generally recommended to use [`prefix(::Model, input)`](@ref) when working with submodels
+    to ensure that the variables in `model` are unique and do not clash with other variables in the
+    parent model or in other submodels.
+
+# Examples
+
+## Simple example
+```jldoctest submodel-to-sampleable; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y)
+            a ~ to_sampleable(demo1(x))
+            return y ~ Uniform(0, a)
+       end;
+```
+
+When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
+```jldoctest submodel-to-sampleable
+julia> vi = VarInfo(demo2(missing, 0.4));
+
+julia> @varname(x) in keys(vi)
+true
+```
+
+Variable `a` is not tracked since it can be computed from the random variable `x` that was
+tracked when running `demo1`:
+```jldoctest submodel-to-sampleable
+julia> @varname(a) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-to-sampleable
+julia> x = vi[@varname(x)];
+
+julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
+true
+```
+
+## With prefixing
+```jldoctest submodel-to-sampleable-prefix; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y, z)
+            a ~ to_sampleable(prefix(demo1(x), :sub1))
+            b ~ to_sampleable(prefix(demo1(y), :sub2))
+            return z ~ Uniform(-a, b)
+       end;
+```
+
+When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
+`sub2.x` will be sampled:
+```jldoctest submodel-to-sampleable-prefix
+julia> vi = VarInfo(demo2(missing, missing, 0.4));
+
+julia> @varname(var"sub1.x") in keys(vi)
+true
+
+julia> @varname(var"sub2.x") in keys(vi)
+true
+```
+
+Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
+`sub2.x` that were tracked when running `demo1`:
+```jldoctest submodel-to-sampleable-prefix
+julia> @varname(a) in keys(vi)
+false
+
+julia> @varname(b) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-to-sampleable-prefix
+julia> sub1_x = vi[@varname(var"sub1.x")];
+
+julia> sub2_x = vi[@varname(var"sub2.x")];
+
+julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
+
+julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
+
+julia> getlogp(vi) ≈ logprior + loglikelihood
+true
+```
+
+## Different ways of setting the prefix
+```jldoctest submodel-to-sampleable-prefix-alts; setup=:(using DynamicPPL, Distributions)
+julia> @model inner() = x ~ Normal()
+inner (generic function with 2 methods)
+
+julia> # When `prefix` is unspecified, no prefix is used.
+       @model submodel_noprefix() = a ~ to_sampleable(inner())
+submodel_noprefix (generic function with 2 methods)
+
+julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
+true
+
+julia> # Using a static string.
+       @model submodel_prefix_string() = a ~ to_sampleable(prefix(inner(), "my prefix"))
+submodel_prefix_string (generic function with 2 methods)
+
+julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+true
+
+julia> # Using string interpolation.
+       @model submodel_prefix_interpolation() = a = to_sampleable(prefix(inner(), "\$(nameof(inner()))"))
+submodel_prefix_interpolation (generic function with 2 methods)
+
+julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+true
+
+julia> # Or using some arbitrary expression.
+       @model submodel_prefix_expr() = a ~ to_sampleable(prefix(inner(), 1 + 2))
+submodel_prefix_expr (generic function with 2 methods)
+
+julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+true
+```
+"""
+to_sampleable(model::Model) = SampleableModelWrapper(model)
 
 """
     need_concretize(expr)
diff --git a/src/contexts.jl b/src/contexts.jl
index 171246610..e38302875 100644
--- a/src/contexts.jl
+++ b/src/contexts.jl
@@ -322,32 +322,6 @@ function prefix(model::Model, ::Val{x}) where {x}
     return contextualize(model, PrefixContext{Symbol(x)}(model.context))
 end
 
-"""
-    @prefix(model, prefix_expr)
-
-Return `model` but with all random variables prefixed by `prefix_expr`.
-
-The result of `prefix_expr` must will be converted to a `Symbol` and used as the prefix.
-
-!!! note
-    This is effectively just a convenience macro for the method [`DynamicPPL.prefix(::Model, x)`](@ref),
-    which automatically converts the result of `prefix_expr` into a `Val` to avoid runtime overheads
-    for static prefixes. For more control over the prefixing, use the method directly.
-
-# Examples
-
-```jldoctest
-julia> @model demo() = x ~ Dirac(1)
-demo (generic function with 2 methods)
-
-julia> rand(@prefix(demo(), :my_prefix))
-(var"my_prefix.x" = 1,)
-```
-"""
-macro prefix(model, prefix_expr)
-    return :($prefix($(esc(model)), $Val{$Symbol($(esc(prefix_expr)))}()))
-end
-
 struct ConditionContext{Values,Ctx<:AbstractContext} <: AbstractContext
     values::Values
     context::Ctx
diff --git a/src/model.jl b/src/model.jl
index 39497d947..75c4b84db 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1250,15 +1250,3 @@ end
 function returned_quantities(model::Model, values, keys)
     return returned_quantities(model, NamedTuple{keys}(values))
 end
-
-"""
-    @returned_quantities(model, input)
-
-Execute `model` and extract the return-values of `model` for `input`.
-
-!!! note
-    This macro is in fact a simple wrapper around the method [`DynamicPPL.returned_quantities`](@ref).
-"""
-macro returned_quantities(model_expr, input_expr)
-    return :($returned_quantities($(esc(model_expr)), $(esc(input_expr))))
-end
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 98dbf6201..3c6178d78 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -6,7 +6,7 @@ Run a Turing `model` nested inside of a Turing model.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use [`@returned_quantities(model)`](@ref) instead.
+    Use `left ~ to_sampleable(model)` instead (see [`to_sampleable`](@ref)).    
 
 # Examples
 
@@ -71,7 +71,7 @@ keeping track of all random variables correctly.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use [`@returned_quantities`](@ref) combined with [`@prefix`](@ref) instead.
+    Use `left ~ to_sampleable(model)` instead (see [`to_sampleable`](@ref)).
 
 # Examples
 ## Example models
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` is deprecated, use `@returned_quantities model` instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
@@ -244,7 +244,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
     return if args_assign === nothing
         ctx = prefix_submodel_context(prefix, ctx)
         quote
-            # Raise deprecation warning to let user know that we recommend using `@returned_quantities`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ to_sampleable(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
@@ -263,7 +263,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
             )
         end
         quote
-            # Raise deprecation warning to let user know that we recommend using `@returned_quantities`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ to_sampleable(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
@@ -273,167 +273,3 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
         end
     end
 end
-
-"""
-    @returned_quantities model
-
-Run `model` nested inside of another model and return the return-values of the `model`.
-
-!!! warning
-    It's generally recommended to use [`@prefix(::Model, input)`](@ref) or
-    [`@prefix(model, prefix_expr)`](@ref) in combination with `@returned_quantities`
-    to ensure that the variables in `model` are unique and do not clash with other variables in the
-    parent model or in other submodels.
-
-# Examples
-
-## Simple example
-```jldoctest submodel-returned-quantities; setup=:(using Distributions)
-julia> @model function demo1(x)
-           x ~ Normal()
-           return 1 + abs(x)
-       end;
-
-julia> @model function demo2(x, y)
-            a = @returned_quantities(demo1(x))
-            return y ~ Uniform(0, a)
-       end;
-```
-
-When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
-```jldoctest submodel-returned-quantities
-julia> vi = VarInfo(demo2(missing, 0.4));
-
-julia> @varname(x) in keys(vi)
-true
-```
-
-Variable `a` is not tracked since it can be computed from the random variable `x` that was
-tracked when running `demo1`:
-```jldoctest submodel-returned-quantities
-julia> @varname(a) in keys(vi)
-false
-```
-
-We can check that the log joint probability of the model accumulated in `vi` is correct:
-
-```jldoctest submodel-returned-quantities
-julia> x = vi[@varname(x)];
-
-julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
-true
-```
-
-## With prefixing
-```jldoctest submodel-returned-quantities-prefix; setup=:(using Distributions)
-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;
-```
-
-When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
-`sub2.x` will be sampled:
-```jldoctest submodel-returned-quantities-prefix
-julia> vi = VarInfo(demo2(missing, missing, 0.4));
-
-julia> @varname(var"sub1.x") in keys(vi)
-true
-
-julia> @varname(var"sub2.x") in keys(vi)
-true
-```
-
-Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
-`sub2.x` that were tracked when running `demo1`:
-```jldoctest submodel-returned-quantities-prefix
-julia> @varname(a) in keys(vi)
-false
-
-julia> @varname(b) in keys(vi)
-false
-```
-
-We can check that the log joint probability of the model accumulated in `vi` is correct:
-
-```jldoctest submodel-returned-quantities-prefix
-julia> sub1_x = vi[@varname(var"sub1.x")];
-
-julia> sub2_x = vi[@varname(var"sub2.x")];
-
-julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
-
-julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
-
-julia> getlogp(vi) ≈ logprior + loglikelihood
-true
-```
-
-## Different ways of setting the prefix
-```jldoctest submodel-returned-quantities-prefix-alts; setup=:(using DynamicPPL, Distributions)
-julia> @model inner() = x ~ Normal()
-inner (generic function with 2 methods)
-
-julia> # When `prefix` is unspecified, no prefix is used.
-       @model submodel_noprefix() = a = @returned_quantities inner()
-submodel_noprefix (generic function with 2 methods)
-
-julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-true
-
-julia> # Using a static string.
-       @model submodel_prefix_string() = a = @returned_quantities @prefix(inner(), "my prefix")
-submodel_prefix_string (generic function with 2 methods)
-
-julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
-true
-
-julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a = @returned_quantities @prefix(inner(), "\$(nameof(inner()))")
-submodel_prefix_interpolation (generic function with 2 methods)
-
-julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
-true
-
-julia> # Or using some arbitrary expression.
-       @model submodel_prefix_expr() = a = @returned_quantities @prefix(inner(), 1 + 2)
-submodel_prefix_expr (generic function with 2 methods)
-
-julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
-true
-```
-"""
-macro returned_quantities(expr)
-    return returned_quantities_expr(expr)
-end
-
-"""
-    returned_quantities_expr(model_expr[, ctx_expr])
-
-Returns an expression that captures the return-values of a model in addition to the varinfo.
-
-!!! warning
-    This is only meant to be used in the body of `@model`.
-"""
-function returned_quantities_expr(model_expr, ctx_expr=esc(:__context__))
-    # The user expects `@returned_quantities ...` to return the
-    # return-value of the `...`, hence we need to capture
-    # the return-value and handle it correctly.
-    @gensym retval
-
-    return quote
-        # Evaluate the model and capture the return values + varinfo.
-        $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
-            $(esc(model_expr)), $(esc(:__varinfo__)), $(ctx_expr)
-        )
-
-        # Return the return-value of the model.
-        $retval
-    end
-end

From 99d99b36b5364af2068ca08002221525a2f00013 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 15 Nov 2024 21:04:37 +0100
Subject: [PATCH 28/63] updated more docstrings

---
 src/submodel_macro.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 3c6178d78..1fd2a1b96 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -25,7 +25,7 @@ julia> @model function demo2(x, y)
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
 ```jldoctest submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -138,7 +138,7 @@ julia> # When `prefix` is unspecified, no prefix is used.
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 true

From 0597b2a10c53507211eb319b086e73f7ff65c322 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 15 Nov 2024 21:19:37 +0100
Subject: [PATCH 29/63] added testing of deprecation warning of `@submodel` +
 replaced some usages in tests (though we don't support some of these so we
 cant' do that yet)

---
 src/submodel_macro.jl |  2 +-
 test/compiler.jl      | 23 ++++++++++++++++++++++-
 test/debug_utils.jl   |  4 ++--
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 1fd2a1b96..707508543 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
diff --git a/test/compiler.jl b/test/compiler.jl
index f2d7e5852..174e9b0c6 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -382,6 +382,27 @@ module Issue537 end
         @test demo2()() == 42
     end
 
+    @testset "@submodel is deprecated" begin
+        @model inner() = x ~ Normal()
+        @model outer() = @submodel x = inner()
+        @test_logs(
+            (
+                :warn,
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead.",
+            ),
+            outer()()
+        )
+
+        @model outer_with_prefix() = @submodel prefix="sub" x = inner()
+        @test_logs(
+            (
+                :warn,
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead.",
+            ),
+            outer_with_prefix()()
+        )
+    end
+
     @testset "submodel" begin
         # No prefix, 1 level.
         @model function demo1(x)
@@ -469,7 +490,7 @@ module Issue537 end
             num_steps = length(y[1])
             num_obs = length(y)
             @inbounds for i in 1:num_obs
-                @submodel prefix = "ar1_$i" x = AR1(num_steps, α, μ, σ)
+                x ~ to_sampleable(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
                 y[i] ~ MvNormal(x, 0.01 * I)
             end
         end
diff --git a/test/debug_utils.jl b/test/debug_utils.jl
index 50bb5d4be..8bc3b50ad 100644
--- a/test/debug_utils.jl
+++ b/test/debug_utils.jl
@@ -45,14 +45,14 @@
         @testset "submodel" begin
             @model ModelInner() = x ~ Normal()
             @model function ModelOuterBroken()
-                @submodel z = ModelInner()
+                z ~ to_sampleable(ModelInner())
                 return x ~ Normal()
             end
             model = ModelOuterBroken()
             @test_throws ErrorException check_model(model; error_on_failure=true)
 
             @model function ModelOuterWorking()
-                @submodel prefix = true z = ModelInner()
+                z = to_sampleable(prefix(ModelInner(), "z"))
                 x ~ Normal()
                 return z
             end

From 0c6bada8cce623d5228d09f3b35bb32cd274b3b4 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 15 Nov 2024 21:23:05 +0100
Subject: [PATCH 30/63] Update test/compiler.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 test/compiler.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/compiler.jl b/test/compiler.jl
index 174e9b0c6..1c662cf21 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -393,7 +393,7 @@ module Issue537 end
             outer()()
         )
 
-        @model outer_with_prefix() = @submodel prefix="sub" x = inner()
+        @model outer_with_prefix() = @submodel prefix = "sub" x = inner()
         @test_logs(
             (
                 :warn,

From 5134ff7f855b69ad9193d616e78b35ff1a9a4877 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 09:27:27 +0100
Subject: [PATCH 31/63] renamed `returned_quantities` to `returned` as
 requested

---
 docs/src/api.md                |  4 ++--
 ext/DynamicPPLMCMCChainsExt.jl |  8 ++++----
 src/DynamicPPL.jl              |  5 +++--
 src/deprecated.jl              |  2 +-
 src/model.jl                   | 18 +++++++++---------
 src/submodel_macro.jl          |  2 +-
 6 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 962939c2c..07465c11f 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -130,10 +130,10 @@ It is possible to manually increase (or decrease) the accumulated log density fr
 @addlogprob!
 ```
 
-Return values of the model function for a collection of samples can be obtained with [`returned_quantities`](@ref).
+Return values of the model function for a collection of samples can be obtained with [`returned(model, chain)`](@ref).
 
 ```@docs
-returned_quantities
+returned(model, chain)
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
index f950de6f1..27e02ddf5 100644
--- a/ext/DynamicPPLMCMCChainsExt.jl
+++ b/ext/DynamicPPLMCMCChainsExt.jl
@@ -43,7 +43,7 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
 end
 
 """
-    returned_quantities(model::Model, chain::MCMCChains.Chains)
+    returned(model::Model, chain::MCMCChains.Chains)
 
 Execute `model` for each of the samples in `chain` and return an array of the values
 returned by the `model` for each sample.
@@ -63,7 +63,7 @@ m = demo(data)
 chain = sample(m, alg, n)
 # To inspect the `interesting_quantity(θ, x)` where `θ` is replaced by samples
 # from the posterior/`chain`:
-returned_quantities(m, chain) # <= results in a `Vector` of returned values
+returned(m, chain) # <= results in a `Vector` of returned values
                                #    from `interesting_quantity(θ, x)`
 ```
 ## Concrete (and simple)
@@ -87,7 +87,7 @@ julia> model = demo(randn(10));
 
 julia> chain = sample(model, MH(), 10);
 
-julia> DynamicPPL.returned_quantities(model, chain)
+julia> returned(model, chain)
 10×1 Array{Tuple{Float64},2}:
  (2.1964758025119338,)
  (2.1964758025119338,)
@@ -101,7 +101,7 @@ julia> DynamicPPL.returned_quantities(model, chain)
  (-0.16489786710222099,)
 ```
 """
-function DynamicPPL.returned_quantities(
+function DynamicPPL.returned(
     model::DynamicPPL.Model, chain_full::MCMCChains.Chains
 )
     chain = MCMCChains.get_sections(chain_full, :parameters)
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 0334bd466..ba12b0fe5 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -122,7 +122,7 @@ export AbstractVarInfo,
     fix,
     unfix,
     prefix,
-    returned_quantities,
+    returned,
     # Convenience macros
     @addlogprob!,
     @submodel,
@@ -132,7 +132,8 @@ export AbstractVarInfo,
     to_sampleable,
     # Deprecated.
     @logprob_str,
-    @prob_str
+    @prob_str,
+    generated_quantities
 
 # Reexport
 using Distributions: loglikelihood
diff --git a/src/deprecated.jl b/src/deprecated.jl
index 32492a9d9..0bcaae9b7 100644
--- a/src/deprecated.jl
+++ b/src/deprecated.jl
@@ -1 +1 @@
-@deprecate generated_quantities returned_quantities
+@deprecate generated_quantities(model, params) returned(model, params)
diff --git a/src/model.jl b/src/model.jl
index 75c4b84db..a05b7c444 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1206,9 +1206,9 @@ function Distributions.loglikelihood(model::Model, chain::AbstractMCMC.AbstractC
 end
 
 """
-    returned_quantities(model::Model, parameters::NamedTuple)
-    returned_quantities(model::Model, values, keys)
-    returned_quantities(model::Model, values, keys)
+    returned(model::Model, parameters::NamedTuple)
+    returned(model::Model, values, keys)
+    returned(model::Model, values, keys)
 
 Execute `model` with variables `keys` set to `values` and return the values returned by the `model`.
 
@@ -1218,7 +1218,7 @@ If a `NamedTuple` is given, `keys=keys(parameters)` and `values=values(parameter
 ```jldoctest
 julia> using DynamicPPL, Distributions
 
-julia> using DynamicPPL: returned_quantities
+julia> using DynamicPPL: returned
 
 julia> @model function demo(xs)
            s ~ InverseGamma(2, 3)
@@ -1235,18 +1235,18 @@ julia> model = demo(randn(10));
 
 julia> parameters = (; s = 1.0, m_shifted=10.0);
 
-julia> returned_quantities(model, parameters)
+julia> returned(model, parameters)
 (0.0,)
 
-julia> returned_quantities(model, values(parameters), keys(parameters))
+julia> returned(model, values(parameters), keys(parameters))
 (0.0,)
 ```
 """
-function returned_quantities(model::Model, parameters::NamedTuple)
+function returned(model::Model, parameters::NamedTuple)
     fixed_model = fix(model, parameters)
     return fixed_model()
 end
 
-function returned_quantities(model::Model, values, keys)
-    return returned_quantities(model, NamedTuple{keys}(values))
+function returned(model::Model, values, keys)
+    return returned(model, NamedTuple{keys}(values))
 end
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 707508543..c2256dc99 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(returned(model))` and `left ~ to_sampleable(returned(prefix(model, ...)))`, respectively, instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)

From 45451f77408bac8721fa16d5bbcdc796ac2ec814 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:22:19 +0100
Subject: [PATCH 32/63] removed redundant `SampleableModelWrapper` in favour of
 `ReturnedModelWrapper` + introduced `rand_like!!` to hide explicit calls to
 `_evaluate!!`

---
 src/compiler.jl                | 194 ++++-----------------------------
 src/context_implementations.jl |  20 +++-
 src/model.jl                   | 194 +++++++++++++++++++++++++++++++++
 3 files changed, 231 insertions(+), 177 deletions(-)

diff --git a/src/compiler.jl b/src/compiler.jl
index f19cd5c75..97b6aa308 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -1,152 +1,5 @@
 const INTERNALNAMES = (:__model__, :__context__, :__varinfo__)
 
-struct SampleableModelWrapper{M}
-    model::M
-end
-
-"""
-    to_sampleable(model::Model)
-
-Return a wrapper around `model` which indicates that this model can only be sampled from.
-
-This is mainly meant to be used on the right-hand side of a `~` operator to indicate that
-the model can be sampled from but not necessarily evaluated for its log density.
-
-!!! warning
-    Note that other operations that one typically associate with expressions of the form `left ~ right`
-    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `to_sampleable`.
-
-!!! warning
-    It's generally recommended to use [`prefix(::Model, input)`](@ref) when working with submodels
-    to ensure that the variables in `model` are unique and do not clash with other variables in the
-    parent model or in other submodels.
-
-# Examples
-
-## Simple example
-```jldoctest submodel-to-sampleable; setup=:(using Distributions)
-julia> @model function demo1(x)
-           x ~ Normal()
-           return 1 + abs(x)
-       end;
-
-julia> @model function demo2(x, y)
-            a ~ to_sampleable(demo1(x))
-            return y ~ Uniform(0, a)
-       end;
-```
-
-When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
-```jldoctest submodel-to-sampleable
-julia> vi = VarInfo(demo2(missing, 0.4));
-
-julia> @varname(x) in keys(vi)
-true
-```
-
-Variable `a` is not tracked since it can be computed from the random variable `x` that was
-tracked when running `demo1`:
-```jldoctest submodel-to-sampleable
-julia> @varname(a) in keys(vi)
-false
-```
-
-We can check that the log joint probability of the model accumulated in `vi` is correct:
-
-```jldoctest submodel-to-sampleable
-julia> x = vi[@varname(x)];
-
-julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
-true
-```
-
-## With prefixing
-```jldoctest submodel-to-sampleable-prefix; setup=:(using Distributions)
-julia> @model function demo1(x)
-           x ~ Normal()
-           return 1 + abs(x)
-       end;
-
-julia> @model function demo2(x, y, z)
-            a ~ to_sampleable(prefix(demo1(x), :sub1))
-            b ~ to_sampleable(prefix(demo1(y), :sub2))
-            return z ~ Uniform(-a, b)
-       end;
-```
-
-When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
-`sub2.x` will be sampled:
-```jldoctest submodel-to-sampleable-prefix
-julia> vi = VarInfo(demo2(missing, missing, 0.4));
-
-julia> @varname(var"sub1.x") in keys(vi)
-true
-
-julia> @varname(var"sub2.x") in keys(vi)
-true
-```
-
-Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
-`sub2.x` that were tracked when running `demo1`:
-```jldoctest submodel-to-sampleable-prefix
-julia> @varname(a) in keys(vi)
-false
-
-julia> @varname(b) in keys(vi)
-false
-```
-
-We can check that the log joint probability of the model accumulated in `vi` is correct:
-
-```jldoctest submodel-to-sampleable-prefix
-julia> sub1_x = vi[@varname(var"sub1.x")];
-
-julia> sub2_x = vi[@varname(var"sub2.x")];
-
-julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
-
-julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
-
-julia> getlogp(vi) ≈ logprior + loglikelihood
-true
-```
-
-## Different ways of setting the prefix
-```jldoctest submodel-to-sampleable-prefix-alts; setup=:(using DynamicPPL, Distributions)
-julia> @model inner() = x ~ Normal()
-inner (generic function with 2 methods)
-
-julia> # When `prefix` is unspecified, no prefix is used.
-       @model submodel_noprefix() = a ~ to_sampleable(inner())
-submodel_noprefix (generic function with 2 methods)
-
-julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-true
-
-julia> # Using a static string.
-       @model submodel_prefix_string() = a ~ to_sampleable(prefix(inner(), "my prefix"))
-submodel_prefix_string (generic function with 2 methods)
-
-julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
-true
-
-julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a = to_sampleable(prefix(inner(), "\$(nameof(inner()))"))
-submodel_prefix_interpolation (generic function with 2 methods)
-
-julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
-true
-
-julia> # Or using some arbitrary expression.
-       @model submodel_prefix_expr() = a ~ to_sampleable(prefix(inner(), 1 + 2))
-submodel_prefix_expr (generic function with 2 methods)
-
-julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
-true
-```
-"""
-to_sampleable(model::Model) = SampleableModelWrapper(model)
-
 """
     need_concretize(expr)
 
@@ -325,6 +178,7 @@ function check_tilde_rhs(@nospecialize(x))
 end
 check_tilde_rhs(x::Distribution) = x
 check_tilde_rhs(x::AbstractArray{<:Distribution}) = x
+check_tilde_rhs(x::ReturnedModelWrapper) = x
 
 """
     unwrap_right_vn(right, vn)
@@ -574,34 +428,28 @@ function generate_tilde(left, right)
     # more selective with our escape. Until that's the case, we remove them all.
     return quote
         $dist = $right
-
-        if $dist isa $(SampleableModelWrapper)
-            $left, __varinfo__ = $(_evaluate!!)($dist.model, __varinfo__, __context__)
-            $left
+        $vn = $(DynamicPPL.resolve_varnames)(
+            $(AbstractPPL.drop_escape(varname(left, need_concretize(left)))), $dist
+        )
+        $isassumption = $(DynamicPPL.isassumption(left, vn))
+        if $(DynamicPPL.isfixed(left, vn))
+            $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
+        elseif $isassumption
+            $(generate_tilde_assume(left, dist, vn))
         else
-            $vn = $(DynamicPPL.resolve_varnames)(
-                $(AbstractPPL.drop_escape(varname(left, need_concretize(left)))), $dist
-            )
-            $isassumption = $(DynamicPPL.isassumption(left, vn))
-            if $(DynamicPPL.isfixed(left, vn))
-                $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
-            elseif $isassumption
-                $(generate_tilde_assume(left, dist, vn))
-            else
-                # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
-                if !$(DynamicPPL.inargnames)($vn, __model__)
-                    $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
-                end
-
-                $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
-                    __context__,
-                    $(DynamicPPL.check_tilde_rhs)($dist),
-                    $(maybe_view(left)),
-                    $vn,
-                    __varinfo__,
-                )
-                $value
+            # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
+            if !$(DynamicPPL.inargnames)($vn, __model__)
+                $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
             end
+
+            $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
+                __context__,
+                $(DynamicPPL.check_tilde_rhs)($dist),
+                $(maybe_view(left)),
+                $vn,
+                __varinfo__,
+            )
+            $value
         end
     end
 end
diff --git a/src/context_implementations.jl b/src/context_implementations.jl
index f3c5171b0..119fa457b 100644
--- a/src/context_implementations.jl
+++ b/src/context_implementations.jl
@@ -141,8 +141,12 @@ By default, calls `tilde_assume(context, right, vn, vi)` and accumulates the log
 probability of `vi` with the returned value.
 """
 function tilde_assume!!(context, right, vn, vi)
-    value, logp, vi = tilde_assume(context, right, vn, vi)
-    return value, acclogp_assume!!(context, vi, logp)
+    return if is_rhs_model(right)
+        rand_like!!(right, context, vi)
+    else
+        value, logp, vi = tilde_assume(context, right, vn, vi)
+        value, acclogp_assume!!(context, vi, logp)
+    end
 end
 
 # observe
@@ -197,6 +201,7 @@ Falls back to `tilde_observe!!(context, right, left, vi)` ignoring the informati
 and indices; if needed, these can be accessed through this function, though.
 """
 function tilde_observe!!(context, right, left, vname, vi)
+    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
     return tilde_observe!!(context, right, left, vi)
 end
 
@@ -210,6 +215,7 @@ By default, calls `tilde_observe(context, right, left, vi)` and accumulates the
 probability of `vi` with the returned value.
 """
 function tilde_observe!!(context, right, left, vi)
+    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
     logp, vi = tilde_observe(context, right, left, vi)
     return left, acclogp_observe!!(context, vi, logp)
 end
@@ -420,8 +426,12 @@ model inputs), accumulate the log probability, and return the sampled value and
 Falls back to `dot_tilde_assume(context, right, left, vn, vi)`.
 """
 function dot_tilde_assume!!(context, right, left, vn, vi)
-    value, logp, vi = dot_tilde_assume(context, right, left, vn, vi)
-    return value, acclogp_assume!!(context, vi, logp), vi
+    return if is_rhs_model(right)
+        rand_like!!(right, context, vi)
+    else
+        value, logp, vi = dot_tilde_assume(context, right, left, vn, vi)
+        value, acclogp_assume!!(context, vi, logp)
+    end
 end
 
 # `dot_assume`
@@ -672,6 +682,7 @@ Falls back to `dot_tilde_observe!!(context, right, left, vi)` ignoring the infor
 name and indices; if needed, these can be accessed through this function, though.
 """
 function dot_tilde_observe!!(context, right, left, vn, vi)
+    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
     return dot_tilde_observe!!(context, right, left, vi)
 end
 
@@ -684,6 +695,7 @@ probability, and return the observed value and updated `vi`.
 Falls back to `dot_tilde_observe(context, right, left, vi)`.
 """
 function dot_tilde_observe!!(context, right, left, vi)
+    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
     logp, vi = dot_tilde_observe(context, right, left, vi)
     return left, acclogp_observe!!(context, vi, logp)
 end
diff --git a/src/model.jl b/src/model.jl
index a05b7c444..3cf722868 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1250,3 +1250,197 @@ end
 function returned(model::Model, values, keys)
     return returned(model, NamedTuple{keys}(values))
 end
+
+"""
+    is_rhs_model(x)
+
+Return `true` if `x` is a model or model wrapper, and `false` otherwise.
+"""
+is_rhs_model(x) = false
+
+
+"""
+    to_sampleable(model)
+
+Return a wrapper around `model` indicating it is sampleable.
+"""
+function to_sampleable end
+
+
+"""
+    ReturnedModelWrapper
+
+A wrapper around a model indicating it is a model over its return values.
+
+This should rarely be constructed explicitly; see [`returned(model)`](@ref) instead.
+"""
+struct ReturnedModelWrapper{M<:Model}
+    model::M
+end
+
+is_rhs_model(::ReturnedModelWrapper) = true
+
+function rand_like!!(model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo)
+    # Return's the value and the (possibly mutated) varinfo.
+    return _evaluate!!(model_wrap.model, varinfo, context)
+end
+
+"""
+    returned(model::Model)
+
+Return a wrapper around `model` which indicates that this model can only be sampled from.
+
+This is mainly meant to be used on the right-hand side of a `~` operator to indicate that
+the model can be sampled from but not necessarily evaluated for its log density.
+
+!!! warning
+    Note that other operations that one typically associate with expressions of the form `left ~ right`
+    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `returned`.
+
+!!! warning
+    It's generally recommended to use [`prefix(::Model, input)`](@ref) when working with submodels
+    to ensure that the variables in `model` are unique and do not clash with other variables in the
+    parent model or in other submodels.
+
+# Examples
+
+## Simple example
+```jldoctest submodel-returned; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y)
+            a ~ returned(demo1(x))
+            return y ~ Uniform(0, a)
+       end;
+```
+
+When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
+```jldoctest submodel-returned
+julia> vi = VarInfo(demo2(missing, 0.4));
+
+julia> @varname(x) in keys(vi)
+true
+```
+
+Variable `a` is not tracked since it can be computed from the random variable `x` that was
+tracked when running `demo1`:
+```jldoctest submodel-returned
+julia> @varname(a) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-returned
+julia> x = vi[@varname(x)];
+
+julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
+true
+```
+
+## With prefixing
+```jldoctest submodel-returned-prefix; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y, z)
+            a ~ returned(prefix(demo1(x), :sub1))
+            b ~ returned(prefix(demo1(y), :sub2))
+            return z ~ Uniform(-a, b)
+       end;
+```
+
+When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
+`sub2.x` will be sampled:
+```jldoctest submodel-returned-prefix
+julia> vi = VarInfo(demo2(missing, missing, 0.4));
+
+julia> @varname(var"sub1.x") in keys(vi)
+true
+
+julia> @varname(var"sub2.x") in keys(vi)
+true
+```
+
+Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
+`sub2.x` that were tracked when running `demo1`:
+```jldoctest submodel-returned-prefix
+julia> @varname(a) in keys(vi)
+false
+
+julia> @varname(b) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-returned-prefix
+julia> sub1_x = vi[@varname(var"sub1.x")];
+
+julia> sub2_x = vi[@varname(var"sub2.x")];
+
+julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
+
+julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
+
+julia> getlogp(vi) ≈ logprior + loglikelihood
+true
+```
+
+## Different ways of setting the prefix
+```jldoctest submodel-returned-prefix-alts; setup=:(using DynamicPPL, Distributions)
+julia> @model inner() = x ~ Normal()
+inner (generic function with 2 methods)
+
+julia> # When `prefix` is unspecified, no prefix is used.
+       @model submodel_noprefix() = a ~ returned(inner())
+submodel_noprefix (generic function with 2 methods)
+
+julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
+true
+
+julia> # Using a static string.
+       @model submodel_prefix_string() = a ~ returned(prefix(inner(), "my prefix"))
+submodel_prefix_string (generic function with 2 methods)
+
+julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+true
+
+julia> # Using string interpolation.
+       @model submodel_prefix_interpolation() = a ~ returned(prefix(inner(), "\$(nameof(inner()))"))
+submodel_prefix_interpolation (generic function with 2 methods)
+
+julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+true
+
+julia> # Or using some arbitrary expression.
+       @model submodel_prefix_expr() = a ~ returned(prefix(inner(), 1 + 2))
+submodel_prefix_expr (generic function with 2 methods)
+
+julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+true
+```
+
+## Usage as likelihood is illegal
+
+Note that it is illegal to use a `returned` model as a likelihood in another model:
+
+```jldoctest submodel-returned-illegal; setup=:(using Distributions)
+julia> @model inner() = x ~ Normal()
+inner (generic function with 2 methods)
+
+julia> @model illegal_likelihood() = a ~ returned(inner())
+illegal_likelihood (generic function with 2 methods)
+
+julia> model = illegal_likelihood() | (a = 1.0,);
+
+julia> model()
+ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
+[...]
+"""
+returned(model::Model) = ReturnedModelWrapper(model)

From c00a9aedf6e3cd435bc69b37661e1c0aa423a825 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:24:22 +0100
Subject: [PATCH 33/63] updated tests + docstrings + warnings to use `returned`

---
 src/submodel_macro.jl | 14 +++++++-------
 test/compiler.jl      |  6 +++---
 test/debug_utils.jl   |  4 ++--
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index c2256dc99..72bab599a 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -6,7 +6,7 @@ Run a Turing `model` nested inside of a Turing model.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use `left ~ to_sampleable(model)` instead (see [`to_sampleable`](@ref)).    
+    Use `left ~ returned(model)` instead (see [`returned(model)`](@ref)).
 
 # Examples
 
@@ -25,7 +25,7 @@ julia> @model function demo2(x, y)
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
 ```jldoctest submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ returned(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -71,7 +71,7 @@ keeping track of all random variables correctly.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use `left ~ to_sampleable(model)` instead (see [`to_sampleable`](@ref)).
+    Use `left ~ returned(model)` instead (see [`returned(model)`](@ref)).
 
 # Examples
 ## Example models
@@ -138,7 +138,7 @@ julia> # When `prefix` is unspecified, no prefix is used.
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-┌ Warning: `@submodel model` is deprecated, use `left ~ to_sampleable(model)` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ returned(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 true
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(returned(model))` and `left ~ to_sampleable(returned(prefix(model, ...)))`, respectively, instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(returned(model))` and `left ~ returned(returned(prefix(model, ...)))`, respectively, instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
@@ -244,7 +244,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
     return if args_assign === nothing
         ctx = prefix_submodel_context(prefix, ctx)
         quote
-            # Raise deprecation warning to let user know that we recommend using `left ~ to_sampleable(model)`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ returned(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
@@ -263,7 +263,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
             )
         end
         quote
-            # Raise deprecation warning to let user know that we recommend using `left ~ to_sampleable(model)`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ returned(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
diff --git a/test/compiler.jl b/test/compiler.jl
index 1c662cf21..687b4622a 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -388,7 +388,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead.",
             ),
             outer()()
         )
@@ -397,7 +397,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_sampleable(model)` and `left ~ to_sampleable(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead.",
             ),
             outer_with_prefix()()
         )
@@ -490,7 +490,7 @@ module Issue537 end
             num_steps = length(y[1])
             num_obs = length(y)
             @inbounds for i in 1:num_obs
-                x ~ to_sampleable(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
+                x ~ returned(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
                 y[i] ~ MvNormal(x, 0.01 * I)
             end
         end
diff --git a/test/debug_utils.jl b/test/debug_utils.jl
index 8bc3b50ad..93095eb8a 100644
--- a/test/debug_utils.jl
+++ b/test/debug_utils.jl
@@ -45,14 +45,14 @@
         @testset "submodel" begin
             @model ModelInner() = x ~ Normal()
             @model function ModelOuterBroken()
-                z ~ to_sampleable(ModelInner())
+                z ~ returned(ModelInner())
                 return x ~ Normal()
             end
             model = ModelOuterBroken()
             @test_throws ErrorException check_model(model; error_on_failure=true)
 
             @model function ModelOuterWorking()
-                z = to_sampleable(prefix(ModelInner(), "z"))
+                z = returned(prefix(ModelInner(), "z"))
                 x ~ Normal()
                 return z
             end

From f0af1d548869d8398358bea78d53204b94ddef39 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:25:36 +0100
Subject: [PATCH 34/63] updated docs

---
 docs/src/api.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 07465c11f..bf6f6ca62 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,13 +14,13 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with `left ~ to_sampleable(model)`.
+One can nest models and call another model inside the model function with `left ~ returned(model)`.
 
 ```@docs
-to_sampleable
+returned(model)
 ```
 
-In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`to_sampleable(model)`](@ref)
+In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`returned(model)`](@ref)
 
 ```@docs
 @submodel

From 1b231a934fd20742f7c8a3e96bffd0c37254a6d0 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:26:20 +0100
Subject: [PATCH 35/63] formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 ext/DynamicPPLMCMCChainsExt.jl |  4 +---
 src/context_implementations.jl | 24 ++++++++++++++++++++----
 src/model.jl                   |  5 +++--
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
index 27e02ddf5..52200171d 100644
--- a/ext/DynamicPPLMCMCChainsExt.jl
+++ b/ext/DynamicPPLMCMCChainsExt.jl
@@ -101,9 +101,7 @@ julia> returned(model, chain)
  (-0.16489786710222099,)
 ```
 """
-function DynamicPPL.returned(
-    model::DynamicPPL.Model, chain_full::MCMCChains.Chains
-)
+function DynamicPPL.returned(model::DynamicPPL.Model, chain_full::MCMCChains.Chains)
     chain = MCMCChains.get_sections(chain_full, :parameters)
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
diff --git a/src/context_implementations.jl b/src/context_implementations.jl
index 119fa457b..6ccd4aa01 100644
--- a/src/context_implementations.jl
+++ b/src/context_implementations.jl
@@ -201,7 +201,11 @@ Falls back to `tilde_observe!!(context, right, left, vi)` ignoring the informati
 and indices; if needed, these can be accessed through this function, though.
 """
 function tilde_observe!!(context, right, left, vname, vi)
-    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
+    is_rhs_model(right) && throw(
+        ArgumentError(
+            "`~` with a model on the right-hand side of an observe statement is not supported",
+        ),
+    )
     return tilde_observe!!(context, right, left, vi)
 end
 
@@ -215,7 +219,11 @@ By default, calls `tilde_observe(context, right, left, vi)` and accumulates the
 probability of `vi` with the returned value.
 """
 function tilde_observe!!(context, right, left, vi)
-    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
+    is_rhs_model(right) && throw(
+        ArgumentError(
+            "`~` with a model on the right-hand side of an observe statement is not supported",
+        ),
+    )
     logp, vi = tilde_observe(context, right, left, vi)
     return left, acclogp_observe!!(context, vi, logp)
 end
@@ -682,7 +690,11 @@ Falls back to `dot_tilde_observe!!(context, right, left, vi)` ignoring the infor
 name and indices; if needed, these can be accessed through this function, though.
 """
 function dot_tilde_observe!!(context, right, left, vn, vi)
-    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
+    is_rhs_model(right) && throw(
+        ArgumentError(
+            "`~` with a model on the right-hand side of an observe statement is not supported",
+        ),
+    )
     return dot_tilde_observe!!(context, right, left, vi)
 end
 
@@ -695,7 +707,11 @@ probability, and return the observed value and updated `vi`.
 Falls back to `dot_tilde_observe(context, right, left, vi)`.
 """
 function dot_tilde_observe!!(context, right, left, vi)
-    is_rhs_model(right) && throw(ArgumentError("`~` with a model on the right-hand side of an observe statement is not supported"))
+    is_rhs_model(right) && throw(
+        ArgumentError(
+            "`~` with a model on the right-hand side of an observe statement is not supported",
+        ),
+    )
     logp, vi = dot_tilde_observe(context, right, left, vi)
     return left, acclogp_observe!!(context, vi, logp)
 end
diff --git a/src/model.jl b/src/model.jl
index 3cf722868..f375a158b 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1266,7 +1266,6 @@ Return a wrapper around `model` indicating it is sampleable.
 """
 function to_sampleable end
 
-
 """
     ReturnedModelWrapper
 
@@ -1280,7 +1279,9 @@ end
 
 is_rhs_model(::ReturnedModelWrapper) = true
 
-function rand_like!!(model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo)
+function rand_like!!(
+    model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo
+)
     # Return's the value and the (possibly mutated) varinfo.
     return _evaluate!!(model_wrap.model, varinfo, context)
 end

From 1faa627d93b66103421de67be90244f7532361a8 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:26:34 +0100
Subject: [PATCH 36/63] Update src/model.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 src/model.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index f375a158b..7bcc441ef 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1258,7 +1258,6 @@ Return `true` if `x` is a model or model wrapper, and `false` otherwise.
 """
 is_rhs_model(x) = false
 
-
 """
     to_sampleable(model)
 

From 92ac6b93b5d674190853685f9f4359f80baa95aa Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:38:15 +0100
Subject: [PATCH 37/63] fix docs

---
 docs/src/api.md   | 4 ++--
 src/DynamicPPL.jl | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index bf6f6ca62..a0e06451d 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -17,7 +17,7 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 One can nest models and call another model inside the model function with `left ~ returned(model)`.
 
 ```@docs
-returned(model)
+returned(::DynamicPPL.Model)
 ```
 
 In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`returned(model)`](@ref)
@@ -133,7 +133,7 @@ It is possible to manually increase (or decrease) the accumulated log density fr
 Return values of the model function for a collection of samples can be obtained with [`returned(model, chain)`](@ref).
 
 ```@docs
-returned(model, chain)
+returned(::DynamicPPL.Model, ::NamedTuple)
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index ba12b0fe5..a1007cbd9 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -129,7 +129,6 @@ export AbstractVarInfo,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,
-    to_sampleable,
     # Deprecated.
     @logprob_str,
     @prob_str,

From b7b2e1d54f4996017cc5c6d1a3194a6be5be98e3 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:46:56 +0100
Subject: [PATCH 38/63] export `to_sampleable` and add to docs

---
 docs/src/api.md   | 7 +++++++
 src/DynamicPPL.jl | 1 +
 2 files changed, 8 insertions(+)

diff --git a/docs/src/api.md b/docs/src/api.md
index a0e06451d..f1d406337 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -20,6 +20,13 @@ One can nest models and call another model inside the model function with `left
 returned(::DynamicPPL.Model)
 ```
 
+Note that a `[returned(::DynamicPPL.Model)](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
+This can be indicated using [`to_sampleable`](@ref) if the user wants to be explicit.
+
+```@docs
+to_sampleable
+```
+
 In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`returned(model)`](@ref)
 
 ```@docs
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index a1007cbd9..4d5aaac7d 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -123,6 +123,7 @@ export AbstractVarInfo,
     unfix,
     prefix,
     returned,
+    to_sampleable,
     # Convenience macros
     @addlogprob!,
     @submodel,

From ed4bb76238b1a2a458c4569ec10a1947700e16ff Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:49:04 +0100
Subject: [PATCH 39/63] fixed typo in warning

---
 src/submodel_macro.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 72bab599a..ff5b07afe 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(returned(model))` and `left ~ returned(returned(prefix(model, ...)))`, respectively, instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)

From 36f02f65ccf24312c558c9a488d5e001abee3c78 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:51:05 +0100
Subject: [PATCH 40/63] removed unnecessary import in docstring

---
 src/model.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 7bcc441ef..a47ff7b88 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1218,8 +1218,6 @@ If a `NamedTuple` is given, `keys=keys(parameters)` and `values=values(parameter
 ```jldoctest
 julia> using DynamicPPL, Distributions
 
-julia> using DynamicPPL: returned
-
 julia> @model function demo(xs)
            s ~ InverseGamma(2, 3)
            m_shifted ~ Normal(10, √s)

From 98538c5d443861689b7268d07c56d6edfa6f7564 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:53:20 +0100
Subject: [PATCH 41/63] added docstring to `rand_like!!`

---
 src/model.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/model.jl b/src/model.jl
index a47ff7b88..73cf83f24 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1276,6 +1276,16 @@ end
 
 is_rhs_model(::ReturnedModelWrapper) = true
 
+"""
+    rand_like!!(model_wrap, context, varinfo)
+    
+Returns a tuple with the first element being the realization and the second the updated varinfo.
+
+# Arguments
+- `model_wrap::ReturnedModelWrapper`: the wrapper of the model to use.
+- `context::AbstractContext`: the context to use for evaluation.
+- `varinfo::AbstractVarInfo`: the varinfo to use for evaluation.
+"""
 function rand_like!!(
     model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo
 )

From d3163067bdd50bde1b8137387f7c8ac84ee1101d Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 10:54:34 +0100
Subject: [PATCH 42/63] fixed docstring for `returned(model)`

---
 src/model.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index 73cf83f24..721d176c1 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1296,7 +1296,7 @@ end
 """
     returned(model::Model)
 
-Return a wrapper around `model` which indicates that this model can only be sampled from.
+Return a `model` wrapper indicating that this a model over the return-values.
 
 This is mainly meant to be used on the right-hand side of a `~` operator to indicate that
 the model can be sampled from but not necessarily evaluated for its log density.

From 0e05901d722b728bf26d3aa4574eb44a378ff457 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 25 Nov 2024 11:29:33 +0100
Subject: [PATCH 43/63] improvements to docstrings thanks to @penelopesym

Co-authored-by: Penelope Yong <penelopeysm@gmail.com>
---
 src/model.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 721d176c1..0d98b560a 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1333,8 +1333,8 @@ julia> @varname(x) in keys(vi)
 true
 ```
 
-Variable `a` is not tracked since it can be computed from the random variable `x` that was
-tracked when running `demo1`:
+The variable `a` is not tracked. However, it will be assigned the return value of `demo1`,
+and can be used in subsequent lines of the model, as shown above.
 ```jldoctest submodel-returned
 julia> @varname(a) in keys(vi)
 false
@@ -1375,8 +1375,8 @@ julia> @varname(var"sub2.x") in keys(vi)
 true
 ```
 
-Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
-`sub2.x` that were tracked when running `demo1`:
+Variables `a` and `b` are not tracked, but are assigned the return values of the respective
+calls to `demo1`:
 ```jldoctest submodel-returned-prefix
 julia> @varname(a) in keys(vi)
 false

From f073b25a6adee34ea14bda74f315bcd5df967665 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 26 Nov 2024 19:36:57 +0100
Subject: [PATCH 44/63] added abstract type `Distributional` and concrete type
 `Sampleable`, in addition to method `to_submodel`

---
 src/DynamicPPL.jl |  1 +
 src/model.jl      | 28 +++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 4d5aaac7d..7e59674a9 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -124,6 +124,7 @@ export AbstractVarInfo,
     prefix,
     returned,
     to_sampleable,
+    to_submodel,
     # Convenience macros
     @addlogprob!,
     @submodel,
diff --git a/src/model.jl b/src/model.jl
index 721d176c1..bf7295605 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1256,12 +1256,28 @@ Return `true` if `x` is a model or model wrapper, and `false` otherwise.
 """
 is_rhs_model(x) = false
 
+"""
+    Distributional
+
+Abstract type for type indicating that something is "distributional".
+"""
+abstract type Distributional end
+
+"""
+    Sampleable{M} <: Distributional
+
+A wrapper around a model indicating it is sampleable.
+"""
+struct Sampleable{M} <: Distributional
+    model::M
+end
+
 """
     to_sampleable(model)
 
 Return a wrapper around `model` indicating it is sampleable.
 """
-function to_sampleable end
+to_sampleable(model) = Sampleable(model)
 
 """
     ReturnedModelWrapper
@@ -1452,3 +1468,13 @@ ERROR: ArgumentError: `~` with a model on the right-hand side of an observe stat
 [...]
 """
 returned(model::Model) = ReturnedModelWrapper(model)
+
+
+"""
+    to_submodel(model::Model)
+
+Return a model wrapper indicating that it is a sampleable model over the return-values.
+
+This can be used on the right-hand side of a `~` operator to include models within other models.
+"""
+to_submodel(model::Model) = to_sampleable(returned(model))

From 2ec03c1a3f456cb5710322dd88e43f8b6fe7ef85 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 26 Nov 2024 19:45:31 +0100
Subject: [PATCH 45/63] replaced usages of `returned` with `to_submodel`

---
 docs/src/api.md       | 48 +++++++++++------------
 src/model.jl          | 88 ++++++++++++++++++++++---------------------
 src/submodel_macro.jl | 16 ++++----
 test/compiler.jl      |  6 +--
 test/debug_utils.jl   |  4 +-
 5 files changed, 82 insertions(+), 80 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index f1d406337..06c94bb44 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,31 +14,6 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with `left ~ returned(model)`.
-
-```@docs
-returned(::DynamicPPL.Model)
-```
-
-Note that a `[returned(::DynamicPPL.Model)](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
-This can be indicated using [`to_sampleable`](@ref) if the user wants to be explicit.
-
-```@docs
-to_sampleable
-```
-
-In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`returned(model)`](@ref)
-
-```@docs
-@submodel
-```
-
-In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
-
-```@docs
-prefix
-```
-
 ### Type
 
 A [`Model`](@ref) can be created by calling the model function, as defined by [`@model`](@ref).
@@ -129,6 +104,29 @@ Similarly, we can [`unfix`](@ref) variables, i.e. return them to their original
 unfix
 ```
 
+## Models within models
+
+One can include models and call another model inside the model function with `left ~ to_submodel(model)`.
+
+```@docs
+to_submodel
+```
+
+Note that a `[returned(::DynamicPPL.Model)](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
+
+In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`to_submodel(model)`](@ref)
+
+```@docs
+@submodel
+```
+
+In the context of including models within models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
+
+```@docs
+prefix
+```
+
+
 ## Utilities
 
 It is possible to manually increase (or decrease) the accumulated log density from within a model function.
diff --git a/src/model.jl b/src/model.jl
index bf7295605..946e4fec9 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1272,6 +1272,8 @@ struct Sampleable{M} <: Distributional
     model::M
 end
 
+is_rhs_model(x::Sampleable) = true
+
 """
     to_sampleable(model)
 
@@ -1279,6 +1281,22 @@ Return a wrapper around `model` indicating it is sampleable.
 """
 to_sampleable(model) = Sampleable(model)
 
+"""
+    rand_like!!(model_wrap, context, varinfo)
+
+Returns a tuple with the first element being the realization and the second the updated varinfo.
+
+# Arguments
+- `model_wrap::ReturnedModelWrapper`: the wrapper of the model to use.
+- `context::AbstractContext`: the context to use for evaluation.
+- `varinfo::AbstractVarInfo`: the varinfo to use for evaluation.
+    """
+function rand_like!!(
+    model_wrap::Sampleable, context::AbstractContext, varinfo::AbstractVarInfo
+)
+    return rand_like!!(model_wrap.model, context, varinfo)
+end
+
 """
     ReturnedModelWrapper
 
@@ -1290,18 +1308,6 @@ struct ReturnedModelWrapper{M<:Model}
     model::M
 end
 
-is_rhs_model(::ReturnedModelWrapper) = true
-
-"""
-    rand_like!!(model_wrap, context, varinfo)
-    
-Returns a tuple with the first element being the realization and the second the updated varinfo.
-
-# Arguments
-- `model_wrap::ReturnedModelWrapper`: the wrapper of the model to use.
-- `context::AbstractContext`: the context to use for evaluation.
-- `varinfo::AbstractVarInfo`: the varinfo to use for evaluation.
-"""
 function rand_like!!(
     model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo
 )
@@ -1313,13 +1319,21 @@ end
     returned(model::Model)
 
 Return a `model` wrapper indicating that this a model over the return-values.
+"""
+returned(model::Model) = ReturnedModelWrapper(model)
+
+
+"""
+    to_submodel(model::Model)
+
+Return a model wrapper indicating that it is a sampleable model over the return-values.
 
 This is mainly meant to be used on the right-hand side of a `~` operator to indicate that
 the model can be sampled from but not necessarily evaluated for its log density.
 
 !!! warning
     Note that other operations that one typically associate with expressions of the form `left ~ right`
-    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `returned`.
+    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `to_submodel`.
 
 !!! warning
     It's generally recommended to use [`prefix(::Model, input)`](@ref) when working with submodels
@@ -1329,20 +1343,20 @@ the model can be sampled from but not necessarily evaluated for its log density.
 # Examples
 
 ## Simple example
-```jldoctest submodel-returned; setup=:(using Distributions)
+```jldoctest submodel-to_submodel; setup=:(using Distributions)
 julia> @model function demo1(x)
            x ~ Normal()
            return 1 + abs(x)
        end;
 
 julia> @model function demo2(x, y)
-            a ~ returned(demo1(x))
+            a ~ to_submodel(demo1(x))
             return y ~ Uniform(0, a)
        end;
 ```
 
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
-```jldoctest submodel-returned
+```jldoctest submodel-to_submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
 
 julia> @varname(x) in keys(vi)
@@ -1351,14 +1365,14 @@ true
 
 Variable `a` is not tracked since it can be computed from the random variable `x` that was
 tracked when running `demo1`:
-```jldoctest submodel-returned
+```jldoctest submodel-to_submodel
 julia> @varname(a) in keys(vi)
 false
 ```
 
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
-```jldoctest submodel-returned
+```jldoctest submodel-to_submodel
 julia> x = vi[@varname(x)];
 
 julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
@@ -1366,22 +1380,22 @@ true
 ```
 
 ## With prefixing
-```jldoctest submodel-returned-prefix; setup=:(using Distributions)
+```jldoctest submodel-to_submodel-prefix; setup=:(using Distributions)
 julia> @model function demo1(x)
            x ~ Normal()
            return 1 + abs(x)
        end;
 
 julia> @model function demo2(x, y, z)
-            a ~ returned(prefix(demo1(x), :sub1))
-            b ~ returned(prefix(demo1(y), :sub2))
+            a ~ to_submodel(prefix(demo1(x), :sub1))
+            b ~ to_submodel(prefix(demo1(y), :sub2))
             return z ~ Uniform(-a, b)
        end;
 ```
 
 When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
 `sub2.x` will be sampled:
-```jldoctest submodel-returned-prefix
+```jldoctest submodel-to_submodel-prefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
 
 julia> @varname(var"sub1.x") in keys(vi)
@@ -1393,7 +1407,7 @@ true
 
 Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
 `sub2.x` that were tracked when running `demo1`:
-```jldoctest submodel-returned-prefix
+```jldoctest submodel-to_submodel-prefix
 julia> @varname(a) in keys(vi)
 false
 
@@ -1403,7 +1417,7 @@ false
 
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
-```jldoctest submodel-returned-prefix
+```jldoctest submodel-to_submodel-prefix
 julia> sub1_x = vi[@varname(var"sub1.x")];
 
 julia> sub2_x = vi[@varname(var"sub2.x")];
@@ -1417,33 +1431,33 @@ true
 ```
 
 ## Different ways of setting the prefix
-```jldoctest submodel-returned-prefix-alts; setup=:(using DynamicPPL, Distributions)
+```jldoctest submodel-to_submodel-prefix-alts; setup=:(using DynamicPPL, Distributions)
 julia> @model inner() = x ~ Normal()
 inner (generic function with 2 methods)
 
 julia> # When `prefix` is unspecified, no prefix is used.
-       @model submodel_noprefix() = a ~ returned(inner())
+       @model submodel_noprefix() = a ~ to_submodel(inner())
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
 true
 
 julia> # Using a static string.
-       @model submodel_prefix_string() = a ~ returned(prefix(inner(), "my prefix"))
+       @model submodel_prefix_string() = a ~ to_submodel(prefix(inner(), "my prefix"))
 submodel_prefix_string (generic function with 2 methods)
 
 julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
 true
 
 julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a ~ returned(prefix(inner(), "\$(nameof(inner()))"))
+       @model submodel_prefix_interpolation() = a ~ to_submodel(prefix(inner(), "\$(nameof(inner()))"))
 submodel_prefix_interpolation (generic function with 2 methods)
 
 julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
 true
 
 julia> # Or using some arbitrary expression.
-       @model submodel_prefix_expr() = a ~ returned(prefix(inner(), 1 + 2))
+       @model submodel_prefix_expr() = a ~ to_submodel(prefix(inner(), 1 + 2))
 submodel_prefix_expr (generic function with 2 methods)
 
 julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
@@ -1452,13 +1466,13 @@ true
 
 ## Usage as likelihood is illegal
 
-Note that it is illegal to use a `returned` model as a likelihood in another model:
+Note that it is illegal to use a `to_submodel` model as a likelihood in another model:
 
-```jldoctest submodel-returned-illegal; setup=:(using Distributions)
+```jldoctest submodel-to_submodel-illegal; setup=:(using Distributions)
 julia> @model inner() = x ~ Normal()
 inner (generic function with 2 methods)
 
-julia> @model illegal_likelihood() = a ~ returned(inner())
+julia> @model illegal_likelihood() = a ~ to_submodel(inner())
 illegal_likelihood (generic function with 2 methods)
 
 julia> model = illegal_likelihood() | (a = 1.0,);
@@ -1467,14 +1481,4 @@ julia> model()
 ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
 [...]
 """
-returned(model::Model) = ReturnedModelWrapper(model)
-
-
-"""
-    to_submodel(model::Model)
-
-Return a model wrapper indicating that it is a sampleable model over the return-values.
-
-This can be used on the right-hand side of a `~` operator to include models within other models.
-"""
 to_submodel(model::Model) = to_sampleable(returned(model))
diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index ff5b07afe..5e1eb987c 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -6,7 +6,7 @@ Run a Turing `model` nested inside of a Turing model.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use `left ~ returned(model)` instead (see [`returned(model)`](@ref)).
+    Use `left ~ to_submodel(model)` instead (see [`to_submodel(model)`](@ref)).
 
 # Examples
 
@@ -25,7 +25,7 @@ julia> @model function demo2(x, y)
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
 ```jldoctest submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `left ~ returned(model)` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ to_submodel(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -71,7 +71,7 @@ keeping track of all random variables correctly.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use `left ~ returned(model)` instead (see [`returned(model)`](@ref)).
+    Use `left ~ to_submodel(model)` instead (see [`to_submodel(model)`](@ref)).
 
 # Examples
 ## Example models
@@ -92,7 +92,7 @@ When we sample from the model `demo2(missing, missing, 0.4)` random variables `s
 `sub2.x` will be sampled:
 ```jldoctest submodelprefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `@returned_quantities model` instead.
+┌ Warning: `@submodel model` is deprecated, use `to_submodel(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -138,7 +138,7 @@ julia> # When `prefix` is unspecified, no prefix is used.
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-┌ Warning: `@submodel model` is deprecated, use `left ~ returned(model)` instead.
+┌ Warning: `@submodel model` is deprecated, use `left ~ to_submodel(model)` instead.
 │   caller = ip:0x0
 └ @ Core :-1
 true
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
@@ -244,7 +244,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
     return if args_assign === nothing
         ctx = prefix_submodel_context(prefix, ctx)
         quote
-            # Raise deprecation warning to let user know that we recommend using `left ~ returned(model)`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ to_submodel(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
@@ -263,7 +263,7 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__))
             )
         end
         quote
-            # Raise deprecation warning to let user know that we recommend using `left ~ returned(model)`.
+            # Raise deprecation warning to let user know that we recommend using `left ~ to_submodel(model)`.
             $(Base.depwarn)(SUBMODEL_DEPWARN_MSG, Symbol("@submodel"))
 
             $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
diff --git a/test/compiler.jl b/test/compiler.jl
index 687b4622a..0537d9a6d 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -388,7 +388,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead.",
             ),
             outer()()
         )
@@ -397,7 +397,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ returned(model)` and `left ~ returned(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead.",
             ),
             outer_with_prefix()()
         )
@@ -490,7 +490,7 @@ module Issue537 end
             num_steps = length(y[1])
             num_obs = length(y)
             @inbounds for i in 1:num_obs
-                x ~ returned(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
+                x ~ to_submodel(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
                 y[i] ~ MvNormal(x, 0.01 * I)
             end
         end
diff --git a/test/debug_utils.jl b/test/debug_utils.jl
index 93095eb8a..e800ebe30 100644
--- a/test/debug_utils.jl
+++ b/test/debug_utils.jl
@@ -45,14 +45,14 @@
         @testset "submodel" begin
             @model ModelInner() = x ~ Normal()
             @model function ModelOuterBroken()
-                z ~ returned(ModelInner())
+                z ~ to_submodel(ModelInner())
                 return x ~ Normal()
             end
             model = ModelOuterBroken()
             @test_throws ErrorException check_model(model; error_on_failure=true)
 
             @model function ModelOuterWorking()
-                z = returned(prefix(ModelInner(), "z"))
+                z = to_submodel(prefix(ModelInner(), "z"))
                 x ~ Normal()
                 return z
             end

From 1f70dfccb5788f0dfcad249bc630f07cd565c0b3 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 26 Nov 2024 19:46:47 +0100
Subject: [PATCH 46/63] formatting

---
 src/model.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index 946e4fec9..774d417d6 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1322,7 +1322,6 @@ Return a `model` wrapper indicating that this a model over the return-values.
 """
 returned(model::Model) = ReturnedModelWrapper(model)
 
-
 """
     to_submodel(model::Model)
 

From 23355ead69823a0d82e15e99bea7ca5da7463bb0 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 27 Nov 2024 07:39:01 +0100
Subject: [PATCH 47/63] Update docs/src/api.md

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 docs/src/api.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 06c94bb44..38b3bf913 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -126,7 +126,6 @@ In the context of including models within models, it's also useful to prefix the
 prefix
 ```
 
-
 ## Utilities
 
 It is possible to manually increase (or decrease) the accumulated log density from within a model function.

From 0e82a6010a8a7bbbc2792529755f24d2b292ad23 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 27 Nov 2024 07:46:33 +0100
Subject: [PATCH 48/63] removed export of `to_sampleable` since it currently
 has no purpose + fixed docs for `returned`

---
 docs/src/api.md   | 7 ++++++-
 src/DynamicPPL.jl | 1 -
 src/model.jl      | 6 ++++--
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 06c94bb44..7448812bf 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -112,7 +112,7 @@ One can include models and call another model inside the model function with `le
 to_submodel
 ```
 
-Note that a `[returned(::DynamicPPL.Model)](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
+Note that a `[to_submodel](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
 
 In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`to_submodel(model)`](@ref)
 
@@ -126,6 +126,11 @@ In the context of including models within models, it's also useful to prefix the
 prefix
 ```
 
+Under the hood, [`to_submodel`](@ref) makes use of the following method to indicate that the model it's wrapping is a model over its return-values rather than something else
+
+```@docs
+returned(::Model)
+```
 
 ## Utilities
 
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 7e59674a9..d4e13d456 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -123,7 +123,6 @@ export AbstractVarInfo,
     unfix,
     prefix,
     returned,
-    to_sampleable,
     to_submodel,
     # Convenience macros
     @addlogprob!,
diff --git a/src/model.jl b/src/model.jl
index 053d4e415..f85892ff3 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1274,6 +1274,7 @@ end
 
 is_rhs_model(x::Sampleable) = true
 
+# TODO: Export this if it end up having a purpose beyond `to_submodel`.
 """
     to_sampleable(model)
 
@@ -1315,10 +1316,11 @@ function rand_like!!(
     return _evaluate!!(model_wrap.model, varinfo, context)
 end
 
+
 """
-    returned(model::Model)
+    returned(model)
 
-Return a `model` wrapper indicating that this a model over the return-values.
+Return a `model` wrapper indicating that it is a model over its return-values.
 """
 returned(model::Model) = ReturnedModelWrapper(model)
 

From b9017c49758f80c99ebf47de103bee769ba5eecd Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 27 Nov 2024 07:47:09 +0100
Subject: [PATCH 49/63] formatting

---
 src/model.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index f85892ff3..404fa2bf7 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1316,7 +1316,6 @@ function rand_like!!(
     return _evaluate!!(model_wrap.model, varinfo, context)
 end
 
-
 """
     returned(model)
 

From 933e4ed5a085a96cef195b2bf4b1440ae7b8976e Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 27 Nov 2024 07:51:20 +0100
Subject: [PATCH 50/63] updated docstring for `condition` and `fix` to not use
 `@submdoel`

---
 src/model.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 404fa2bf7..eb44f3880 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -223,14 +223,14 @@ true
 ## Nested models
 
 `condition` of course also supports the use of nested models through
-the use of [`@submodel`](@ref).
+the use of [`to_submodel`](@ref).
 
 ```jldoctest condition
 julia> @model demo_inner() = m ~ Normal()
 demo_inner (generic function with 2 methods)
 
 julia> @model function demo_outer()
-           @submodel m = demo_inner()
+           m ~ to_submodel(demo_inner())
            return m
        end
 demo_outer (generic function with 2 methods)
@@ -250,7 +250,7 @@ But one needs to be careful when prefixing variables in the nested models:
 
 ```jldoctest condition
 julia> @model function demo_outer_prefix()
-           @submodel prefix="inner" m = demo_inner()
+           m ~ to_submodel(prefix(demo_inner(), "inner"))
            return m
        end
 demo_outer_prefix (generic function with 2 methods)
@@ -578,14 +578,14 @@ true
 ## Nested models
 
 `fix` of course also supports the use of nested models through
-the use of [`@submodel`](@ref).
+the use of [`to_submodel`](@ref).
 
 ```jldoctest fix
 julia> @model demo_inner() = m ~ Normal()
 demo_inner (generic function with 2 methods)
 
 julia> @model function demo_outer()
-           @submodel m = demo_inner()
+           m ~ to_submodel(demo_inner())
            return m
        end
 demo_outer (generic function with 2 methods)
@@ -605,7 +605,7 @@ But one needs to be careful when prefixing variables in the nested models:
 
 ```jldoctest fix
 julia> @model function demo_outer_prefix()
-           @submodel prefix="inner" m = demo_inner()
+           m ~ to_submodel(prefix(demo_inner(), "inner"))
            return m
        end
 demo_outer_prefix (generic function with 2 methods)

From 4fc7b76b566b7b1b08afb6dac8eb57df8f31ba96 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 27 Nov 2024 08:38:16 +0100
Subject: [PATCH 51/63] added `check_tilde_rhs` for `Sampleable`

---
 src/compiler.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/compiler.jl b/src/compiler.jl
index 97b6aa308..bb066e431 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -179,6 +179,7 @@ end
 check_tilde_rhs(x::Distribution) = x
 check_tilde_rhs(x::AbstractArray{<:Distribution}) = x
 check_tilde_rhs(x::ReturnedModelWrapper) = x
+check_tilde_rhs(x::Sampleable) = Sampleable(check_tilde_rhs(x.model))
 
 """
     unwrap_right_vn(right, vn)

From b42168731bcf4c905bfdaba020863c0a66dfbab7 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 28 Nov 2024 14:36:16 +0100
Subject: [PATCH 52/63] let the field of sampleable determine whether it works
 or not

---
 src/model.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/model.jl b/src/model.jl
index eb44f3880..3897531a0 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1272,7 +1272,7 @@ struct Sampleable{M} <: Distributional
     model::M
 end
 
-is_rhs_model(x::Sampleable) = true
+is_rhs_model(x::Sampleable) = is_rhs_model(x.model)
 
 # TODO: Export this if it end up having a purpose beyond `to_submodel`.
 """
@@ -1309,6 +1309,8 @@ struct ReturnedModelWrapper{M<:Model}
     model::M
 end
 
+is_rhs_model(::ReturnedModelWrapper) = true
+
 function rand_like!!(
     model_wrap::ReturnedModelWrapper, context::AbstractContext, varinfo::AbstractVarInfo
 )

From 3c204d9296f8de88d6da4b12ab20cd89e2f11959 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 28 Nov 2024 21:02:41 +0100
Subject: [PATCH 53/63] add automatic prefixing of submodels + remove support
 for dot-tilde since this is ambigious in this case

---
 src/context_implementations.jl | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/context_implementations.jl b/src/context_implementations.jl
index 175b7fe81..a1118172b 100644
--- a/src/context_implementations.jl
+++ b/src/context_implementations.jl
@@ -104,7 +104,8 @@ probability of `vi` with the returned value.
 """
 function tilde_assume!!(context, right, vn, vi)
     return if is_rhs_model(right)
-        rand_like!!(right, context, vi)
+        # Prefix the variables using the `vn`.
+        rand_like!!(right, prefix(context Symbol(vn)), vi)
     else
         value, logp, vi = tilde_assume(context, right, vn, vi)
         value, acclogp_assume!!(context, vi, logp)
@@ -335,12 +336,13 @@ model inputs), accumulate the log probability, and return the sampled value and
 Falls back to `dot_tilde_assume(context, right, left, vn, vi)`.
 """
 function dot_tilde_assume!!(context, right, left, vn, vi)
-    return if is_rhs_model(right)
-        rand_like!!(right, context, vi)
-    else
-        value, logp, vi = dot_tilde_assume(context, right, left, vn, vi)
-        value, acclogp_assume!!(context, vi, logp)
-    end
+    is_rhs_model(right) && throw(
+        ArgumentError(
+            "`.~` with a model on the right-hand side is not supported; please use `~`",
+        ),
+    )
+    value, logp, vi = dot_tilde_assume(context, right, left, vn, vi)
+    value, acclogp_assume!!(context, vi, logp)
 end
 
 # `dot_assume`

From 044f6c328dd88dfe0425bf03e93256be1243c02c Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 28 Nov 2024 21:30:36 +0100
Subject: [PATCH 54/63] added automatic prefixing for sub-models involved in
 `~` statements

---
 src/compiler.jl                |   5 +-
 src/context_implementations.jl |   2 +-
 src/model.jl                   | 104 +++++++++++++++++----------------
 3 files changed, 58 insertions(+), 53 deletions(-)

diff --git a/src/compiler.jl b/src/compiler.jl
index bb066e431..c67da6f95 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -179,7 +179,10 @@ end
 check_tilde_rhs(x::Distribution) = x
 check_tilde_rhs(x::AbstractArray{<:Distribution}) = x
 check_tilde_rhs(x::ReturnedModelWrapper) = x
-check_tilde_rhs(x::Sampleable) = Sampleable(check_tilde_rhs(x.model))
+function check_tilde_rhs(x::Sampleable{<:Any,AutoPrefix}) where {AutoPrefix}
+    model = check_tilde_rhs(x.model)
+    return Sampleable{typeof(model),AutoPrefix}(model)
+end
 
 """
     unwrap_right_vn(right, vn)
diff --git a/src/context_implementations.jl b/src/context_implementations.jl
index a1118172b..f876d441d 100644
--- a/src/context_implementations.jl
+++ b/src/context_implementations.jl
@@ -105,7 +105,7 @@ probability of `vi` with the returned value.
 function tilde_assume!!(context, right, vn, vi)
     return if is_rhs_model(right)
         # Prefix the variables using the `vn`.
-        rand_like!!(right, prefix(context Symbol(vn)), vi)
+        rand_like!!(right, should_auto_prefix(right) ? PrefixContext{Symbol(vn)}(context) : context, vi)
     else
         value, logp, vi = tilde_assume(context, right, vn, vi)
         value, acclogp_assume!!(context, vi, logp)
diff --git a/src/model.jl b/src/model.jl
index 3897531a0..af8d8af1c 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1263,24 +1263,43 @@ Abstract type for type indicating that something is "distributional".
 """
 abstract type Distributional end
 
+"""
+    should_auto_prefix(distributional)
+
+Return `true` if the `distributional` should use automatic prefixing, and `false` otherwise.
+"""
+function should_auto_prefix end
+
+"""
+    is_rhs_model(x)
+
+Return `true` if the `distributional` is a model, and `false` otherwise.
+"""
+function is_rhs_model end
+
 """
     Sampleable{M} <: Distributional
 
 A wrapper around a model indicating it is sampleable.
 """
-struct Sampleable{M} <: Distributional
+struct Sampleable{M,AutoPrefix} <: Distributional
     model::M
 end
 
+should_auto_prefix(::Sampleable{<:Any,AutoPrefix}) where {AutoPrefix} = AutoPrefix
 is_rhs_model(x::Sampleable) = is_rhs_model(x.model)
 
 # TODO: Export this if it end up having a purpose beyond `to_submodel`.
 """
-    to_sampleable(model)
+    to_sampleable(model[, auto_prefix])
 
 Return a wrapper around `model` indicating it is sampleable.
+
+# Arguments
+- `model::Model`: the model to wrap.
+- `auto_prefix::Bool`: whether to prefix the variables in the model. Default: `true`.
 """
-to_sampleable(model) = Sampleable(model)
+to_sampleable(model, auto_prefix::Bool=true) = Sampleable{typeof(model),auto_prefix}(model)
 
 """
     rand_like!!(model_wrap, context, varinfo)
@@ -1326,7 +1345,7 @@ Return a `model` wrapper indicating that it is a model over its return-values.
 returned(model::Model) = ReturnedModelWrapper(model)
 
 """
-    to_submodel(model::Model)
+    to_submodel(model::Model[, auto_prefix::Bool])
 
 Return a model wrapper indicating that it is a sampleable model over the return-values.
 
@@ -1338,9 +1357,13 @@ the model can be sampled from but not necessarily evaluated for its log density.
     such as [`condition`](@ref) or [`fix`](@ref), will also not work with `to_submodel`.
 
 !!! warning
-    It's generally recommended to use [`prefix(::Model, input)`](@ref) when working with submodels
-    to ensure that the variables in `model` are unique and do not clash with other variables in the
-    parent model or in other submodels.
+    To avoid variable names clashing between models, it is recommend leave argument `auto_prefix` equal to `true`.
+    If one does not use automatic prefixing, then it's recommended to use [`prefix(::Model, input)`](@ref) explicitly.
+
+# Arguments
+- `model::Model`: the model to wrap.
+- `auto_prefix::Bool`: whether to automatically prefix the variables in the model using the left-hand
+    side of the `~` statement. Default: `true`.
 
 # Examples
 
@@ -1361,7 +1384,7 @@ When we sample from the model `demo2(missing, 0.4)` random variable `x` will be
 ```jldoctest submodel-to_submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
 
-julia> @varname(x) in keys(vi)
+julia> @varname(var\"a.x\") in keys(vi)
 true
 ```
 
@@ -1375,29 +1398,42 @@ false
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
 ```jldoctest submodel-to_submodel
-julia> x = vi[@varname(x)];
+julia> x = vi[@varname(var\"a.x\")];
 
 julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
 true
 ```
 
-## With prefixing
+## Without automatic prefixing
+As mentioned earlier, by default, the `auto_prefix` argument specifies whether to automatically
+prefix the variables in the submodel. If `auto_prefix=false`, then the variables in the submodel
+will not be prefixed.
 ```jldoctest submodel-to_submodel-prefix; setup=:(using Distributions)
 julia> @model function demo1(x)
            x ~ Normal()
            return 1 + abs(x)
        end;
 
+julia> @model function demo2_no_prefix(x, z)
+            a ~ to_submodel(demo1(x), false)
+            return z ~ Uniform(-a, 1)
+       end;
+
+julia> vi = VarInfo(demo2_no_prefix(missing, 0.4));
+
+julia> @varname(x) in keys(vi)  # here we just use `x` instead of `a.x`
+true
+```
+However, not using prefixing is generally not recommended as it can lead to variable name clashes
+unless one is careful. For example, if we're re-using the same model twice in a model, not using prefixing
+will lead to variable name clashes: However, one can manually prefix using the [`prefix(::Model, input)`](@ref):
+```jldoctest submodel-to_submodel-prefix
 julia> @model function demo2(x, y, z)
-            a ~ to_submodel(prefix(demo1(x), :sub1))
-            b ~ to_submodel(prefix(demo1(y), :sub2))
+            a ~ to_submodel(prefix(demo1(x), :sub1), false)
+            b ~ to_submodel(prefix(demo1(y), :sub2), false)
             return z ~ Uniform(-a, b)
        end;
-```
 
-When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
-`sub2.x` will be sampled:
-```jldoctest submodel-to_submodel-prefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
 
 julia> @varname(var"sub1.x") in keys(vi)
@@ -1432,40 +1468,6 @@ julia> getlogp(vi) ≈ logprior + loglikelihood
 true
 ```
 
-## Different ways of setting the prefix
-```jldoctest submodel-to_submodel-prefix-alts; setup=:(using DynamicPPL, Distributions)
-julia> @model inner() = x ~ Normal()
-inner (generic function with 2 methods)
-
-julia> # When `prefix` is unspecified, no prefix is used.
-       @model submodel_noprefix() = a ~ to_submodel(inner())
-submodel_noprefix (generic function with 2 methods)
-
-julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-true
-
-julia> # Using a static string.
-       @model submodel_prefix_string() = a ~ to_submodel(prefix(inner(), "my prefix"))
-submodel_prefix_string (generic function with 2 methods)
-
-julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
-true
-
-julia> # Using string interpolation.
-       @model submodel_prefix_interpolation() = a ~ to_submodel(prefix(inner(), "\$(nameof(inner()))"))
-submodel_prefix_interpolation (generic function with 2 methods)
-
-julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
-true
-
-julia> # Or using some arbitrary expression.
-       @model submodel_prefix_expr() = a ~ to_submodel(prefix(inner(), 1 + 2))
-submodel_prefix_expr (generic function with 2 methods)
-
-julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
-true
-```
-
 ## Usage as likelihood is illegal
 
 Note that it is illegal to use a `to_submodel` model as a likelihood in another model:
@@ -1483,4 +1485,4 @@ julia> model()
 ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
 [...]
 """
-to_submodel(model::Model) = to_sampleable(returned(model))
+to_submodel(model::Model, auto_prefix::Bool=true) = to_sampleable(returned(model), auto_prefix)

From f7162962691c05722c1926f79206977644299639 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 28 Nov 2024 21:35:47 +0100
Subject: [PATCH 55/63] updated depwarn for `@submodel` and tests

---
 src/submodel_macro.jl | 10 +++++-----
 test/compiler.jl      |  6 +++---
 test/debug_utils.jl   |  6 ++++--
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 5e1eb987c..2258b880a 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -6,7 +6,7 @@ Run a Turing `model` nested inside of a Turing model.
 
 !!! warning
     This is deprecated and will be removed in a future release.
-    Use `left ~ to_submodel(model)` instead (see [`to_submodel(model)`](@ref)).
+    Use `left ~ to_submodel(model)` instead (see [`to_submodel`](@ref)).
 
 # Examples
 
@@ -25,7 +25,7 @@ julia> @model function demo2(x, y)
 When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
 ```jldoctest submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `left ~ to_submodel(model)` instead.
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -92,7 +92,7 @@ When we sample from the model `demo2(missing, missing, 0.4)` random variables `s
 `sub2.x` will be sampled:
 ```jldoctest submodelprefix
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
-┌ Warning: `@submodel model` is deprecated, use `to_submodel(model)` instead.
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
 
@@ -138,7 +138,7 @@ julia> # When `prefix` is unspecified, no prefix is used.
 submodel_noprefix (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
-┌ Warning: `@submodel model` is deprecated, use `left ~ to_submodel(model)` instead.
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
 true
@@ -224,7 +224,7 @@ function prefix_submodel_context(prefix::Bool, ctx)
     return ctx
 end
 
-const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead."
+const SUBMODEL_DEPWARN_MSG = "`@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax."
 
 function submodel(prefix_expr, expr, ctx=esc(:__context__))
     prefix_left, prefix = getargs_assignment(prefix_expr)
diff --git a/test/compiler.jl b/test/compiler.jl
index 0537d9a6d..977c1156c 100644
--- a/test/compiler.jl
+++ b/test/compiler.jl
@@ -388,7 +388,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.",
             ),
             outer()()
         )
@@ -397,7 +397,7 @@ module Issue537 end
         @test_logs(
             (
                 :warn,
-                "`@submodel model` and `@submodel prefix=... model` are deprecated, use `left ~ to_submodel(model)` and `left ~ to_submodel(prefix(model, ...))`, respectively, instead.",
+                "`@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.",
             ),
             outer_with_prefix()()
         )
@@ -490,7 +490,7 @@ module Issue537 end
             num_steps = length(y[1])
             num_obs = length(y)
             @inbounds for i in 1:num_obs
-                x ~ to_submodel(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"))
+                x ~ to_submodel(prefix(AR1(num_steps, α, μ, σ), "ar1_$i"), false)
                 y[i] ~ MvNormal(x, 0.01 * I)
             end
         end
diff --git a/test/debug_utils.jl b/test/debug_utils.jl
index 091810971..dfa46affc 100644
--- a/test/debug_utils.jl
+++ b/test/debug_utils.jl
@@ -45,14 +45,16 @@
         @testset "submodel" begin
             @model ModelInner() = x ~ Normal()
             @model function ModelOuterBroken()
-                z ~ to_submodel(ModelInner())
+                # Without automatic prefixing => `x` s used twice.
+                z ~ to_submodel(ModelInner(), false)
                 return x ~ Normal()
             end
             model = ModelOuterBroken()
             @test_throws ErrorException check_model(model; error_on_failure=true)
 
             @model function ModelOuterWorking()
-                z = to_submodel(prefix(ModelInner(), "z"))
+                # With automatic prefixing => `x` is not duplicated.
+                z ~ to_submodel(ModelInner())
                 x ~ Normal()
                 return z
             end

From 76aebc525199cb6964fbd4e1779fab977d0bd466 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 28 Nov 2024 21:36:56 +0100
Subject: [PATCH 56/63] formatting

---
 src/context_implementations.jl | 10 +++++++---
 src/model.jl                   |  3 ++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/context_implementations.jl b/src/context_implementations.jl
index f876d441d..2c3173adc 100644
--- a/src/context_implementations.jl
+++ b/src/context_implementations.jl
@@ -105,7 +105,11 @@ probability of `vi` with the returned value.
 function tilde_assume!!(context, right, vn, vi)
     return if is_rhs_model(right)
         # Prefix the variables using the `vn`.
-        rand_like!!(right, should_auto_prefix(right) ? PrefixContext{Symbol(vn)}(context) : context, vi)
+        rand_like!!(
+            right,
+            should_auto_prefix(right) ? PrefixContext{Symbol(vn)}(context) : context,
+            vi,
+        )
     else
         value, logp, vi = tilde_assume(context, right, vn, vi)
         value, acclogp_assume!!(context, vi, logp)
@@ -338,11 +342,11 @@ Falls back to `dot_tilde_assume(context, right, left, vn, vi)`.
 function dot_tilde_assume!!(context, right, left, vn, vi)
     is_rhs_model(right) && throw(
         ArgumentError(
-            "`.~` with a model on the right-hand side is not supported; please use `~`",
+            "`.~` with a model on the right-hand side is not supported; please use `~`"
         ),
     )
     value, logp, vi = dot_tilde_assume(context, right, left, vn, vi)
-    value, acclogp_assume!!(context, vi, logp)
+    return value, acclogp_assume!!(context, vi, logp)
 end
 
 # `dot_assume`
diff --git a/src/model.jl b/src/model.jl
index af8d8af1c..f7db9ed61 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1485,4 +1485,5 @@ julia> model()
 ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
 [...]
 """
-to_submodel(model::Model, auto_prefix::Bool=true) = to_sampleable(returned(model), auto_prefix)
+to_submodel(model::Model, auto_prefix::Bool=true) =
+    to_sampleable(returned(model), auto_prefix)

From c150a8778a6d1646bbb0b4364716c82e921bf3fe Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 29 Nov 2024 09:52:14 +0100
Subject: [PATCH 57/63] updated docstrings

---
 src/model.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index f7db9ed61..ef4b2546e 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -250,7 +250,7 @@ But one needs to be careful when prefixing variables in the nested models:
 
 ```jldoctest condition
 julia> @model function demo_outer_prefix()
-           m ~ to_submodel(prefix(demo_inner(), "inner"))
+           m ~ to_submodel(prefix(demo_inner(), "inner"), false)
            return m
        end
 demo_outer_prefix (generic function with 2 methods)
@@ -605,7 +605,7 @@ But one needs to be careful when prefixing variables in the nested models:
 
 ```jldoctest fix
 julia> @model function demo_outer_prefix()
-           m ~ to_submodel(prefix(demo_inner(), "inner"))
+           m ~ to_submodel(prefix(demo_inner(), "inner"), false)
            return m
        end
 demo_outer_prefix (generic function with 2 methods)

From ecb473796033a5c64eed6cd77be27cd0b83dc656 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 29 Nov 2024 14:06:38 +0100
Subject: [PATCH 58/63] updated docs

---
 src/model.jl | 109 ++++++++++++---------------------------------------
 1 file changed, 24 insertions(+), 85 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index ef4b2546e..2278943b3 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -230,8 +230,9 @@ julia> @model demo_inner() = m ~ Normal()
 demo_inner (generic function with 2 methods)
 
 julia> @model function demo_outer()
-           m ~ to_submodel(demo_inner())
-           return m
+           # By default, `to_submodel` prefixes the variables using the left-hand side of `~`.
+           inner ~ to_submodel(demo_inner())
+           return inner
        end
 demo_outer (generic function with 2 methods)
 
@@ -240,63 +241,28 @@ julia> model = demo_outer();
 julia> model() ≠ 1.0
 true
 
-julia> conditioned_model = model | (m = 1.0, );
+julia> # To condition the variable inside `demo_inner` we need to refer to it as `inner.m`.
+       conditioned_model = model | (var"inner.m" = 1.0, );
 
 julia> conditioned_model()
 1.0
-```
-
-But one needs to be careful when prefixing variables in the nested models:
-
-```jldoctest condition
-julia> @model function demo_outer_prefix()
-           m ~ to_submodel(prefix(demo_inner(), "inner"), false)
-           return m
-       end
-demo_outer_prefix (generic function with 2 methods)
-
-julia> # (×) This doesn't work now!
-       conditioned_model = demo_outer_prefix() | (m = 1.0, );
-
-julia> conditioned_model() == 1.0
-false
 
-julia> # (✓) `m` in `demo_inner` is referred to as `inner.m` internally, so we do:
-       conditioned_model = demo_outer_prefix() | (var"inner.m" = 1.0, );
+julia> # However, it's not possible to condition `inner` directly.
+       conditioned_model_fail = model | (inner = 1.0, );
 
-julia> conditioned_model()
-1.0
-
-julia> # Note that the above `var"..."` is just standard Julia syntax:
-       keys((var"inner.m" = 1.0, ))
-(Symbol("inner.m"),)
+julia> conditioned_model_fail()
+ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
+[...]
 ```
 
 And similarly when using `Dict`:
 
 ```jldoctest condition
-julia> conditioned_model_dict = demo_outer_prefix() | (@varname(var"inner.m") => 1.0);
+julia> conditioned_model_dict = model | (@varname(var"inner.m") => 1.0);
 
 julia> conditioned_model_dict()
 1.0
 ```
-
-The difference is maybe more obvious once we look at how these different
-in their trace/`VarInfo`:
-
-```jldoctest condition
-julia> keys(VarInfo(demo_outer()))
-1-element Vector{VarName{:m, typeof(identity)}}:
- m
-
-julia> keys(VarInfo(demo_outer_prefix()))
-1-element Vector{VarName{Symbol("inner.m"), typeof(identity)}}:
- inner.m
-```
-
-From this we can tell what the correct way to condition `m` within `demo_inner`
-is in the two different models.
-
 """
 AbstractPPL.condition(model::Model; values...) = condition(model, NamedTuple(values))
 function AbstractPPL.condition(model::Model, value, values...)
@@ -578,15 +544,15 @@ true
 ## Nested models
 
 `fix` of course also supports the use of nested models through
-the use of [`to_submodel`](@ref).
+the use of [`to_submodel`](@ref), similar to [`condition`](@ref).
 
 ```jldoctest fix
 julia> @model demo_inner() = m ~ Normal()
 demo_inner (generic function with 2 methods)
 
 julia> @model function demo_outer()
-           m ~ to_submodel(demo_inner())
-           return m
+           inner ~ to_submodel(demo_inner())
+           return inner
        end
 demo_outer (generic function with 2 methods)
 
@@ -595,63 +561,36 @@ julia> model = demo_outer();
 julia> model() ≠ 1.0
 true
 
-julia> fixed_model = model | (m = 1.0, );
+julia> fixed_model = fix(model, var"inner.m" = 1.0, );
 
 julia> fixed_model()
 1.0
 ```
 
-But one needs to be careful when prefixing variables in the nested models:
-
-```jldoctest fix
-julia> @model function demo_outer_prefix()
-           m ~ to_submodel(prefix(demo_inner(), "inner"), false)
-           return m
-       end
-demo_outer_prefix (generic function with 2 methods)
+However, unlike [`condition`](@ref), `fix` can also be used to fix the
+return-value of the submodel:
 
-julia> # (×) This doesn't work now!
-       fixed_model = demo_outer_prefix() | (m = 1.0, );
-
-julia> fixed_model() == 1.0
-false
-
-julia> # (✓) `m` in `demo_inner` is referred to as `inner.m` internally, so we do:
-       fixed_model = demo_outer_prefix() | (var"inner.m" = 1.0, );
+```julia
+julia> fixed_model = fix(model, inner = 2.0,);
 
 julia> fixed_model()
-1.0
-
-julia> # Note that the above `var"..."` is just standard Julia syntax:
-       keys((var"inner.m" = 1.0, ))
-(Symbol("inner.m"),)
+2.0
 ```
 
 And similarly when using `Dict`:
 
 ```jldoctest fix
-julia> fixed_model_dict = demo_outer_prefix() | (@varname(var"inner.m") => 1.0);
+julia> fixed_model_dict = fix(model, @varname(var"inner.m") => 1.0);
 
 julia> fixed_model_dict()
 1.0
-```
 
-The difference is maybe more obvious once we look at how these different
-in their trace/`VarInfo`:
+julia> fixed_model_dict = fix(model, @varname(inner) => 2.0);
 
-```jldoctest fix
-julia> keys(VarInfo(demo_outer()))
-1-element Vector{VarName{:m, typeof(identity)}}:
- m
-
-julia> keys(VarInfo(demo_outer_prefix()))
-1-element Vector{VarName{Symbol("inner.m"), typeof(identity)}}:
- inner.m
+julia> fixed_model_dict()
+2.0
 ```
 
-From this we can tell what the correct way to fix `m` within `demo_inner`
-is in the two different models.
-
 ## Difference from `condition`
 
 A very similar functionality is also provided by [`condition`](@ref) which,

From ed0cec3324e83ecc87d16c7f39af12c5a5fd68f7 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 2 Dec 2024 13:40:39 +0000
Subject: [PATCH 59/63] added more depwarns to the doctests to see if that
 helps (though I don't understand why this is needed for Documenter.jl)

---
 src/submodel_macro.jl | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 2258b880a..1a91991a1 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -148,6 +148,9 @@ julia> # Explicitely don't use any prefix.
 submodel_prefix_false (generic function with 2 methods)
 
 julia> @varname(x) in keys(VarInfo(submodel_prefix_false()))
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # Automatically determined from `a`.
@@ -155,6 +158,9 @@ julia> # Automatically determined from `a`.
 submodel_prefix_true (generic function with 2 methods)
 
 julia> @varname(var"a.x") in keys(VarInfo(submodel_prefix_true()))
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # Using a static string.
@@ -162,6 +168,9 @@ julia> # Using a static string.
 submodel_prefix_string (generic function with 2 methods)
 
 julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # Using string interpolation.
@@ -169,6 +178,9 @@ julia> # Using string interpolation.
 submodel_prefix_interpolation (generic function with 2 methods)
 
 julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # Or using some arbitrary expression.

From 13a2bf7898587681e414c39e61a15c8bf5b07f06 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 2 Dec 2024 13:41:19 +0000
Subject: [PATCH 60/63] forgot one

---
 src/submodel_macro.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl
index 1a91991a1..e5a8e0617 100644
--- a/src/submodel_macro.jl
+++ b/src/submodel_macro.jl
@@ -188,6 +188,9 @@ julia> # Or using some arbitrary expression.
 submodel_prefix_expr (generic function with 2 methods)
 
 julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
+│   caller = ip:0x0
+└ @ Core :-1
 true
 
 julia> # (×) Automatic prefixing without a left-hand side expression does not work!

From f94e07ac37281b38624e93dcf76d6121a4c46bf0 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Mon, 2 Dec 2024 14:41:52 +0000
Subject: [PATCH 61/63] replaced usage of `generated_quantities` with
 `returned`

---
 test/ext/DynamicPPLMCMCChainsExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/ext/DynamicPPLMCMCChainsExt.jl b/test/ext/DynamicPPLMCMCChainsExt.jl
index c19bf6f2d..e117c4bbc 100644
--- a/test/ext/DynamicPPLMCMCChainsExt.jl
+++ b/test/ext/DynamicPPLMCMCChainsExt.jl
@@ -3,7 +3,7 @@
     model = demo()
 
     chain = MCMCChains.Chains(randn(1000, 2, 1), [:x, :y], Dict(:internals => [:y]))
-    chain_generated = @test_nowarn generated_quantities(model, chain)
+    chain_generated = @test_nowarn returned(model, chain)
     @test size(chain_generated) == (1000, 1)
     @test mean(chain_generated) ≈ 0 atol = 0.1
 end

From d03eb4ccc2b839b6413e18ff7dc67828a89a9efb Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 3 Dec 2024 17:05:33 +0000
Subject: [PATCH 62/63] foxed docstring for `to_submodel`

---
 src/model.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 2278943b3..c0489b089 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -1292,8 +1292,8 @@ This is mainly meant to be used on the right-hand side of a `~` operator to indi
 the model can be sampled from but not necessarily evaluated for its log density.
 
 !!! warning
-    Note that other operations that one typically associate with expressions of the form `left ~ right`
-    such as [`condition`](@ref) or [`fix`](@ref), will also not work with `to_submodel`.
+    Note that some other operations that one typically associate with expressions of the form
+    `left ~ right` such as [`condition`](@ref), will also not work with `to_submodel`.
 
 !!! warning
     To avoid variable names clashing between models, it is recommend leave argument `auto_prefix` equal to `true`.

From 341b6b86eb43b7a3157098fcff384c1486c62a44 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 3 Dec 2024 17:06:09 +0000
Subject: [PATCH 63/63] patch version bump

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 301e485f7..fafeec2ba 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.31.0"
+version = "0.31.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"