From 1d6b305f79c73963a3c8803b11db15286df516f5 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 6 Dec 2021 13:54:23 +0000 Subject: [PATCH 01/10] added prefix keyword argument to submodel-macro --- src/submodel_macro.jl | 120 +++++++++++++++++++++++++++++++++++++----- test/compiler.jl | 2 +- 2 files changed, 108 insertions(+), 14 deletions(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index 54477a05d..e15c10afb 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -1,5 +1,6 @@ """ @submodel model + @submodel ... = model Run a Turing `model` nested inside of a Turing model. @@ -44,22 +45,32 @@ true ``` """ macro submodel(expr) - return submodel(expr) + return submodel(:(prefix=false), expr) end """ - @submodel prefix model + @submodel prefix=... model + @submodel prefix=... ... = model Run a Turing `model` nested inside of a Turing model and add "`prefix`." as a prefix to all random variables inside of the `model`. +Valid expressions for `prefix=...` are: +- `prefix=false`: no prefix is used. +- `prefix=true`: _attempt_ to automatically determine the prefix from the left-hand side + `... = model` by first converting into a `VarName`, and then calling `Symbol` on this. +- `prefix="my prefix"`: prefix is taken to be the static string "my prefix". +- `prefix=expression`: `expression` is evaluated at runtime, resulting in + the prefix `Symbol(expression)`. Note that this also includes string-interpolation, + e.g. `prefix="x[\$i]"` as it requires runtime information. + The prefix makes it possible to run the same Turing model multiple times while keeping track of all random variables correctly. The return value can be assigned to a variable. # Examples - +## Example models ```jldoctest submodelprefix; setup=:(using Distributions) julia> @model function demo1(x) x ~ Normal() @@ -67,8 +78,8 @@ julia> @model function demo1(x) end; julia> @model function demo2(x, y, z) - @submodel sub1 a = demo1(x) - @submodel sub2 b = demo1(y) + @submodel prefix="sub1" a = demo1(x) + @submodel prefix="sub2" b = demo1(y) return z ~ Uniform(-a, b) end; ``` @@ -109,27 +120,110 @@ 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 +```julia +julia> # Don't use any prefix. + @model outer() = @submodel prefix=false a = inner(); + +julia> @varname(x) in keys(VarInfo(outer())) +true + +julia> # Automatically determined from `a`. + @model outer() = @submodel prefix=true a = inner(); + +julia> @varname(var"a.x") in keys(VarInfo(outer())) +true + +julia> # Using a static string. + @model outer() = @submodel prefix="my prefix" a = inner(); + +julia> @varname(var"my prefix.x") in keys(VarInfo(outer())) +true + +julia> # Using string interpolation. + @model outer() = @submodel prefix="$(inner().name)" a = inner(); + +julia> @varname(var"inner.x") in keys(VarInfo(outer())) +true + +julia> # Or using some arbitrary expression. + @model outer() = @submodel prefix=1 + 2 a = inner(); + +julia> @varname(var"3.x") in keys(VarInfo(outer())) +true +``` + +# Notes +- The choice `prefix=expression` means that the prefixing will incur a runtime cost. + This is also the case for `prefix=true`, depending on whether the expression on the + the right-hand side of `... = model` requires runtime-information or not, e.g. + `x = model` will result in the _static_ prefix `x`, while `x[i] = model` will be + resolved at runtime. """ -macro submodel(prefix, expr) - ctx = :(PrefixContext{$(esc(Meta.quot(prefix)))}($(esc(:__context__)))) - return submodel(expr, ctx) +macro submodel(prefix_expr, expr) + return submodel(prefix_expr, expr, esc(:__context__)) +end + +# Automatic prefixing. +function prefix_submodel_context(prefix::Bool, left::Symbol, ctx) + return prefix ? prefix_submodel_context(string(left), ctx) : ctx +end + +function prefix_submodel_context(prefix::Bool, left::Expr, ctx) + return prefix ? prefix_submodel_context(varname(left), ctx) : ctx +end + +# Manual prefixing. +prefix_submodel_context(prefix, left, ctx) = prefix_submodel_context(prefix, ctx) +prefix_submodel_context(prefix, ctx) = ctx + +function prefix_submodel_context(prefix::Union{Symbol,Expr}, ctx) + # E.g. `prefix="asd[$i]"` or `prefix=asd` with `asd` to be evaluated. + return :($(DynamicPPL.PrefixContext){$(Symbol)($(esc(prefix)))}($ctx)) end -function submodel(expr, ctx=esc(:__context__)) +function prefix_submodel_context(prefix::AbstractString, ctx) + # E.g. `prefix="asd"`. + return :($(DynamicPPL.PrefixContext){$(esc(Meta.quot(Symbol(prefix))))}($ctx)) +end + +function prefix_submodel_context(prefix::Bool, ctx) + if prefix + error("cannot automatically prefix with no left-hand side") + end + + return ctx +end + +function submodel(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 + + # `prefix=false` => don't prefix, i.e. do nothing to `ctx`. + # `prefix=true` => automatically determine prefix. + # `prefix=...` => use it. args_assign = getargs_assignment(expr) return if args_assign === nothing + ctx = prefix_submodel_context(prefix, ctx) # In this case we only want to get the `__varinfo__`. quote $(esc(:__varinfo__)) = last( - _evaluate!!($(esc(expr)), $(esc(:__varinfo__)), $(ctx)) + $(DynamicPPL._evaluate!!)($(esc(expr)), $(esc(:__varinfo__)), $(ctx)) ) end else - # Here we also want the return-variable. - # TODO: Should we prefix by `L` by default? L, R = args_assign + # Now that we have `L` and `R`, we can prefix automagically. + try + ctx = prefix_submodel_context(prefix, L, ctx) + catch e + error("failed to determine prefix from $(L); please specify prefix using the `@submodel prefix=\"your prefix\" ...` syntax") + end quote - $(esc(L)), $(esc(:__varinfo__)) = _evaluate!!( + $(esc(L)), $(esc(:__varinfo__)) = $(DynamicPPL._evaluate!!)( $(esc(R)), $(esc(:__varinfo__)), $(ctx) ) end diff --git a/test/compiler.jl b/test/compiler.jl index 80fab9669..402a10065 100644 --- a/test/compiler.jl +++ b/test/compiler.jl @@ -455,7 +455,7 @@ end num_steps = length(y[1]) num_obs = length(y) @inbounds for i in 1:num_obs - @submodel $(Symbol("ar1_$i")) x = AR1(num_steps, α, μ, σ) + @submodel prefix="ar1_$i" x = AR1(num_steps, α, μ, σ) y[i] ~ MvNormal(x, 0.1) end end From 83e552702fe80a58879e56af571887c53b1f47e8 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 6 Dec 2021 14:02:17 +0000 Subject: [PATCH 02/10] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/submodel_macro.jl | 7 ++++--- test/compiler.jl | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index e15c10afb..3a0adc649 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -45,7 +45,7 @@ true ``` """ macro submodel(expr) - return submodel(:(prefix=false), expr) + return submodel(:(prefix = false), expr) end """ @@ -201,7 +201,6 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__)) if prefix_left !== :prefix error("$(prefix_left) is not a valid kwarg") end - # `prefix=false` => don't prefix, i.e. do nothing to `ctx`. # `prefix=true` => automatically determine prefix. # `prefix=...` => use it. @@ -220,7 +219,9 @@ function submodel(prefix_expr, expr, ctx=esc(:__context__)) try ctx = prefix_submodel_context(prefix, L, ctx) catch e - error("failed to determine prefix from $(L); please specify prefix using the `@submodel prefix=\"your prefix\" ...` syntax") + error( + "failed to determine prefix from $(L); please specify prefix using the `@submodel prefix=\"your prefix\" ...` syntax", + ) end quote $(esc(L)), $(esc(:__varinfo__)) = $(DynamicPPL._evaluate!!)( diff --git a/test/compiler.jl b/test/compiler.jl index 402a10065..f5da23839 100644 --- a/test/compiler.jl +++ b/test/compiler.jl @@ -455,7 +455,7 @@ 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, α, μ, σ) + @submodel prefix = "ar1_$i" x = AR1(num_steps, α, μ, σ) y[i] ~ MvNormal(x, 0.1) end end From f1ec6a088b99386fc368f0fa82db2ee45f17056d Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 6 Dec 2021 21:21:33 +0000 Subject: [PATCH 03/10] converted example in docs into test --- src/submodel_macro.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index 3a0adc649..cf995db21 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -122,7 +122,9 @@ true ``` ## Different ways of setting the prefix -```julia +```jldoctest submodel-prefix; setup=:(using DynamicPPL, Distributions) +julia> @model inner() = x ~ Normal(); + julia> # Don't use any prefix. @model outer() = @submodel prefix=false a = inner(); From 3e2ac78856864f9ce00f558fa33b91738e273814 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 7 Dec 2021 23:59:27 +0000 Subject: [PATCH 04/10] fixed docstring --- 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 cf995db21..18f4a237d 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -122,7 +122,7 @@ true ``` ## Different ways of setting the prefix -```jldoctest submodel-prefix; setup=:(using DynamicPPL, Distributions) +```jldoctest submodel-prefix-alternatives; setup=:(using DynamicPPL, Distributions) julia> @model inner() = x ~ Normal(); julia> # Don't use any prefix. @@ -144,7 +144,7 @@ julia> @varname(var"my prefix.x") in keys(VarInfo(outer())) true julia> # Using string interpolation. - @model outer() = @submodel prefix="$(inner().name)" a = inner(); + @model outer() = @submodel prefix="\$(inner().name)" a = inner(); julia> @varname(var"inner.x") in keys(VarInfo(outer())) true From 1b5fc00a2fe299845187abe044df910a7fa18d6a Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 15:10:17 +0000 Subject: [PATCH 05/10] Apply suggestions from code review Co-authored-by: Philipp Gabler --- src/submodel_macro.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index 18f4a237d..10b224f7a 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -125,7 +125,13 @@ true ```jldoctest submodel-prefix-alternatives; setup=:(using DynamicPPL, Distributions) julia> @model inner() = x ~ Normal(); -julia> # Don't use any prefix. +julia> # When `prefix` is unspecified, no prefix is used. + @model outer() = @submodel a = inner(); + +julia> @varname(x) in keys(VarInfo(outer())) +true + +julia> # Explicitely don't use any prefix. @model outer() = @submodel prefix=false a = inner(); julia> @varname(x) in keys(VarInfo(outer())) From 18d07e82ea719afc6062f0bb10d367e5aa4a0412 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 15:27:05 +0000 Subject: [PATCH 06/10] removed redundant prefix_submodel_context def and added another example to docstring --- src/submodel_macro.jl | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index 10b224f7a..4e9a6dfa6 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -159,7 +159,11 @@ julia> # Or using some arbitrary expression. @model outer() = @submodel prefix=1 + 2 a = inner(); julia> @varname(var"3.x") in keys(VarInfo(outer())) -true + +julia> # (×) Automatic prefixing without a left-hand side expression does not work! + @model outer() = @submodel prefix=true inner(); +ERROR: LoadError: LoadError: cannot automatically prefix with no left-hand side +[...] ``` # Notes @@ -175,7 +179,7 @@ end # Automatic prefixing. function prefix_submodel_context(prefix::Bool, left::Symbol, ctx) - return prefix ? prefix_submodel_context(string(left), ctx) : ctx + return prefix ? prefix_submodel_context(left, ctx) : ctx end function prefix_submodel_context(prefix::Bool, left::Expr, ctx) @@ -184,14 +188,12 @@ end # Manual prefixing. prefix_submodel_context(prefix, left, ctx) = prefix_submodel_context(prefix, ctx) -prefix_submodel_context(prefix, ctx) = ctx - -function prefix_submodel_context(prefix::Union{Symbol,Expr}, ctx) +function prefix_submodel_context(prefix, ctx) # E.g. `prefix="asd[$i]"` or `prefix=asd` with `asd` to be evaluated. return :($(DynamicPPL.PrefixContext){$(Symbol)($(esc(prefix)))}($ctx)) end -function prefix_submodel_context(prefix::AbstractString, ctx) +function prefix_submodel_context(prefix::Union{AbstractString,Symbol}, ctx) # E.g. `prefix="asd"`. return :($(DynamicPPL.PrefixContext){$(esc(Meta.quot(Symbol(prefix))))}($ctx)) end From 89d0b38857f7036789536f76fcaf6eeb146ef921 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 15:51:59 +0000 Subject: [PATCH 07/10] fixed doctests --- src/model.jl | 2 +- src/submodel_macro.jl | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/model.jl b/src/model.jl index db7e1e265..e020b8c0b 100644 --- a/src/model.jl +++ b/src/model.jl @@ -215,7 +215,7 @@ But one needs to be careful when prefixing variables in the nested models: ```jldoctest condition julia> @model function demo_outer_prefix() - @submodel inner m = demo_inner() + @submodel prefix="inner" m = demo_inner() return m end demo_outer_prefix (generic function with 2 methods) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index 4e9a6dfa6..b6db5de80 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -159,6 +159,7 @@ julia> # Or using some arbitrary expression. @model outer() = @submodel prefix=1 + 2 a = inner(); julia> @varname(var"3.x") in keys(VarInfo(outer())) +true julia> # (×) Automatic prefixing without a left-hand side expression does not work! @model outer() = @submodel prefix=true inner(); From dc9564fe56fd07e82514aba72e1f689414f1d700 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 16:16:14 +0000 Subject: [PATCH 08/10] attempt at fixing doctests --- src/submodel_macro.jl | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/submodel_macro.jl b/src/submodel_macro.jl index b6db5de80..59e2eb94b 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -123,46 +123,53 @@ true ## Different ways of setting the prefix ```jldoctest submodel-prefix-alternatives; setup=:(using DynamicPPL, Distributions) -julia> @model inner() = x ~ Normal(); +julia> @model inner() = x ~ Normal() +inner (generic function with 2 methods) julia> # When `prefix` is unspecified, no prefix is used. - @model outer() = @submodel a = inner(); - + @model outer() = @submodel a = inner() +outer (generic function with 2 methods) + julia> @varname(x) in keys(VarInfo(outer())) true julia> # Explicitely don't use any prefix. - @model outer() = @submodel prefix=false a = inner(); + @model outer() = @submodel prefix=false a = inner() +outer (generic function with 2 methods) julia> @varname(x) in keys(VarInfo(outer())) true julia> # Automatically determined from `a`. - @model outer() = @submodel prefix=true a = inner(); + @model outer() = @submodel prefix=true a = inner() +outer (generic function with 2 methods) julia> @varname(var"a.x") in keys(VarInfo(outer())) true julia> # Using a static string. - @model outer() = @submodel prefix="my prefix" a = inner(); + @model outer() = @submodel prefix="my prefix" a = inner() +outer (generic function with 2 methods) julia> @varname(var"my prefix.x") in keys(VarInfo(outer())) true julia> # Using string interpolation. - @model outer() = @submodel prefix="\$(inner().name)" a = inner(); + @model outer() = @submodel prefix="\$(inner().name)" a = inner() +outer (generic function with 2 methods) julia> @varname(var"inner.x") in keys(VarInfo(outer())) -true +false julia> # Or using some arbitrary expression. - @model outer() = @submodel prefix=1 + 2 a = inner(); + @model outer() = @submodel prefix=1 + 2 a = inner() +outer (generic function with 2 methods) julia> @varname(var"3.x") in keys(VarInfo(outer())) true julia> # (×) Automatic prefixing without a left-hand side expression does not work! - @model outer() = @submodel prefix=true inner(); + @model outer() = @submodel prefix=true inner() ERROR: LoadError: LoadError: cannot automatically prefix with no left-hand side [...] ``` From 80b5c6b3e792cad8cbbef1e46adcbcf720712479 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 16:20:53 +0000 Subject: [PATCH 09/10] another attempt at fixing doctests --- 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 59e2eb94b..ae48302e6 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -170,7 +170,7 @@ true julia> # (×) Automatic prefixing without a left-hand side expression does not work! @model outer() = @submodel prefix=true inner() -ERROR: LoadError: LoadError: cannot automatically prefix with no left-hand side +ERROR: LoadError: cannot automatically prefix with no left-hand side [...] ``` From 2980cb0f18e48f60acb28e8975068726ab3a614c Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 8 Dec 2021 19:24:59 +0000 Subject: [PATCH 10/10] had a typo in docstring --- 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 ae48302e6..009a12c7b 100644 --- a/src/submodel_macro.jl +++ b/src/submodel_macro.jl @@ -159,7 +159,7 @@ julia> # Using string interpolation. outer (generic function with 2 methods) julia> @varname(var"inner.x") in keys(VarInfo(outer())) -false +true julia> # Or using some arbitrary expression. @model outer() = @submodel prefix=1 + 2 a = inner()