Reopening of previous PR (#309)

What we need before merge:
1. [x] Updated benchmarks.
2. [ ] Resolve comments.

Co-authored-by: Hong Ge <hg344@cam.ac.uk>

Author: torfjelde

.github/workflows/CI.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       matrix:
         version:
-          - '1.3' # minimum supported version
+          # - '1.3' # minimum supported version
           - '1' # current stable version
         os:
           - ubuntu-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.16.2"
+version = "0.17.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

docs/Project.toml

@@ -1,9 +1,11 @@
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
 [compat]
 Distributions = "0.25"
 Documenter = "0.27"
+Setfield = "0.7.1, 0.8"
 StableRNGs = "1"

docs/make.jl

@@ -8,7 +8,7 @@ makedocs(;
     sitename="DynamicPPL",
     format=Documenter.HTML(),
     modules=[DynamicPPL],
-    pages=["Home" => "index.md"],
+    pages=["Home" => "index.md", "TestUtils" => "test_utils.md"],
     strict=true,
     checkdocs=:exports,
     doctestfilters=[

docs/src/test_utils.md

@@ -0,0 +1,5 @@
+# DynamicPPL.TestUtils
+
+```@autodocs
+Modules = [DynamicPPL.TestUtils]
+```

src/DynamicPPL.jl

@@ -10,6 +10,7 @@ using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
 using BangBang: BangBang
+using Setfield: Setfield
 
 using Setfield: Setfield
 using BangBang: BangBang
@@ -31,15 +32,23 @@ import Base:
     keys,
     haskey
 
+using BangBang: push!!, empty!!, setindex!!
+
 # VarInfo
 export AbstractVarInfo,
     VarInfo,
     UntypedVarInfo,
     TypedVarInfo,
+    SimpleVarInfo,
+    push!!,
+    empty!!,
     getlogp,
     setlogp!,
     acclogp!,
     resetlogp!,
+    setlogp!!,
+    acclogp!!,
+    resetlogp!!,
     get_num_produce,
     set_num_produce!,
     reset_num_produce!,
@@ -139,13 +148,32 @@ include("distribution_wrappers.jl")
 include("contexts.jl")
 include("varinfo.jl")
 include("threadsafe.jl")
+include("simple_varinfo.jl")
 include("context_implementations.jl")
 include("compiler.jl")
 include("prob_macro.jl")
 include("compat/ad.jl")
 include("loglikelihoods.jl")
 include("submodel_macro.jl")
-
 include("test_utils.jl")
 
+# Deprecations
+@deprecate empty!(vi::VarInfo) empty!!(vi::VarInfo)
+@deprecate push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution) push!!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution
+)
+@deprecate push!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, sampler::AbstractSampler
+) push!!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, sampler::AbstractSampler)
+@deprecate push!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Selector) push!!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Selector
+)
+@deprecate push!(
+    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Set{Selector}
+) push!!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Set{Selector})
+
+@deprecate setlogp!(vi, logp) setlogp!!(vi, logp)
+@deprecate acclogp!(vi, logp) acclogp!!(vi, logp)
+@deprecate resetlogp!(vi) resetlogp!!(vi)
+
 end # module

src/compat/ad.jl

@@ -1,5 +1,5 @@
 # See https://github.com/TuringLang/Turing.jl/issues/1199
-ChainRulesCore.@non_differentiable push!(
+ChainRulesCore.@non_differentiable push!!(
     vi::VarInfo, vn::VarName, r, dist::Distribution, gidset::Set{Selector}
 )
 
@@ -16,7 +16,7 @@ ZygoteRules.@adjoint function dot_observe(
 )
     function dot_observe_fallback(spl, dists, value, vi)
         increment_num_produce!(vi)
-        return sum(map(Distributions.loglikelihood, dists, value))
+        return sum(map(Distributions.loglikelihood, dists, value)), vi
     end
     return ZygoteRules.pullback(__context__, dot_observe_fallback, spl, dists, value, vi)
 end

src/compiler.jl

@@ -355,10 +355,12 @@ end
 
 function generate_tilde_literal(left, right)
     # If the LHS is a literal, it is always an observation
+    @gensym value
     return quote
-        $(DynamicPPL.tilde_observe!)(
+        $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
             __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
         )
+        $value
     end
 end
 
@@ -373,7 +375,7 @@ function generate_tilde(left, right)
 
     # Otherwise it is determined by the model or its value,
     # if the LHS represents an observation
-    @gensym vn isassumption
+    @gensym vn isassumption value
 
     # HACK: Usage of `drop_escape` is unfortunate. It's a consequence of the fact
     # that in DynamicPPL we the entire function body. Instead we should be
@@ -389,32 +391,38 @@ function generate_tilde(left, right)
                 $left = $(DynamicPPL.getvalue_nested)(__context__, $vn)
             end
 
-            $(DynamicPPL.tilde_observe!)(
+            $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),
                 $(maybe_view(left)),
                 $vn,
                 __varinfo__,
             )
+            $value
         end
     end
 end
 
 function generate_tilde_assume(left, right, vn)
-    expr = :(
-        $left = $(DynamicPPL.tilde_assume!)(
+    # HACK: Because the Setfield.jl macro does not support assignment
+    # with multiple arguments on the LHS, we need to capture the return-values
+    # and then update the LHS variables one by one.
+    @gensym value
+    expr = :($left = $value)
+    if left isa Expr
+        expr = AbstractPPL.drop_escape(
+            Setfield.setmacro(BangBang.prefermutation, expr; overwrite=true)
+        )
+    end
+
+    return quote
+        $value, __varinfo__ = $(DynamicPPL.tilde_assume!!)(
             __context__,
             $(DynamicPPL.unwrap_right_vn)($(DynamicPPL.check_tilde_rhs)($right), $vn)...,
             __varinfo__,
         )
-    )
-
-    return if left isa Expr
-        AbstractPPL.drop_escape(
-            Setfield.setmacro(BangBang.prefermutation, expr; overwrite=true)
-        )
-    else
-        return expr
+        $expr
+        $value
     end
 end
 
@@ -428,7 +436,7 @@ function generate_dot_tilde(left, right)
 
     # Otherwise it is determined by the model or its value,
     # if the LHS represents an observation
-    @gensym vn isassumption
+    @gensym vn isassumption value
     return quote
         $vn = $(AbstractPPL.drop_escape(varname(left)))
         $isassumption = $(DynamicPPL.isassumption(left))
@@ -440,13 +448,14 @@ function generate_dot_tilde(left, right)
                 $left .= $(DynamicPPL.getvalue_nested)(__context__, $vn)
             end
 
-            $(DynamicPPL.dot_tilde_observe!)(
+            $value, __varinfo__ = $(DynamicPPL.dot_tilde_observe!!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),
                 $(maybe_view(left)),
                 $vn,
                 __varinfo__,
             )
+            $value
         end
     end
 end
@@ -455,15 +464,82 @@ function generate_dot_tilde_assume(left, right, vn)
     # We don't need to use `Setfield.@set` here since
     # `.=` is always going to be inplace + needs `left` to
     # be something that supports `.=`.
-    return :(
-        $left .= $(DynamicPPL.dot_tilde_assume!)(
+    @gensym value
+    return quote
+        $value, __varinfo__ = $(DynamicPPL.dot_tilde_assume!!)(
             __context__,
             $(DynamicPPL.unwrap_right_left_vns)(
                 $(DynamicPPL.check_tilde_rhs)($right), $(maybe_view(left)), $vn
             )...,
             __varinfo__,
         )
-    )
+        $left .= $value
+        $value
+    end
+end
+
+# Note that we cannot use `MacroTools.isdef` because
+# of https://github.com/FluxML/MacroTools.jl/issues/154.
+"""
+    isfuncdef(expr)
+
+Return `true` if `expr` is any form of function definition, and `false` otherwise.
+"""
+function isfuncdef(e::Expr)
+    return if Meta.isexpr(e, :function)
+        # Classic `function f(...)`
+        true
+    elseif Meta.isexpr(e, :->)
+        # Anonymous functions/lambdas, e.g. `do` blocks or `->` defs.
+        true
+    elseif Meta.isexpr(e, :(=)) && Meta.isexpr(e.args[1], :call)
+        # Short function defs, e.g. `f(args...) = ...`.
+        true
+    else
+        false
+    end
+end
+
+"""
+    replace_returns(expr)
+
+Return `Expr` with all `return ...` statements replaced with
+`return ..., DynamicPPL.return_values(__varinfo__)`.
+
+Note that this method will _not_ replace `return` statements within function
+definitions. This is checked using [`isfuncdef`](@ref).
+"""
+replace_returns(e) = e
+function replace_returns(e::Expr)
+    if isfuncdef(e)
+        return e
+    end
+
+    if Meta.isexpr(e, :return)
+        # NOTE: `return` always has an argument. In the case of
+        # an empty `return`, the lowered expression will be `return nothing`.
+        # Hence we don't need any special handling for empty returns.
+        retval_expr = if length(e.args) > 1
+            Expr(:tuple, e.args...)
+        else
+            e.args[1]
+        end
+
+        return :(return ($retval_expr, __varinfo__))
+    end
+
+    return Expr(e.head, map(replace_returns, e.args)...)
+end
+
+# If it's just a symbol, e.g. `f(x) = 1`, then we make it `f(x) = return 1`.
+make_returns_explicit!(body) = Expr(:return, body)
+function make_returns_explicit!(body::Expr)
+    # If the last statement is a return-statement, we don't do anything.
+    # Otherwise we replace the last statement with a `return` statement.
+    if !Meta.isexpr(body.args[end], :return)
+        body.args[end] = Expr(:return, body.args[end])
+    end
+    return body
 end
 
 const FloatOrArrayType = Type{<:Union{AbstractFloat,AbstractArray}}
@@ -496,10 +572,14 @@ function build_output(modelinfo, linenumbernode)
     # Replace the user-provided function body with the version created by DynamicPPL.
     # We use `MacroTools.@q begin ... end` instead of regular `quote ... end` to ensure
     # that no new `LineNumberNode`s are added apart from the reference `linenumbernode`
-    # to the call site
+    # to the call site.
+    # NOTE: We need to replace statements of the form `return ...` with
+    # `return (..., __varinfo__)` to ensure that the second
+    # element in the returned value is always the most up-to-date `__varinfo__`.
+    # See the docstrings of `replace_returns` for more info.
     evaluatordef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        $(modelinfo[:body])
+        $(replace_returns(make_returns_explicit!(modelinfo[:body])))
     end
 
     ## Build the model function.