Merge remote-tracking branch 'origin/main' into breaking

Author: mhauru

.github/workflows/Benchmarking.yml

@@ -59,14 +59,14 @@ jobs:
           echo "DPPL_COMMIT_URL=$COMMIT_URL" >> $GITHUB_ENV
 
       - name: Find Existing Comment
-        uses: peter-evans/find-comment@v3
+        uses: peter-evans/find-comment@v4
         id: find_comment
         with:
           issue-number: ${{ github.event.pull_request.number }}
           comment-author: github-actions[bot]
 
       - name: Post Benchmark Results as PR Comment
-        uses: peter-evans/create-or-update-comment@v4
+        uses: peter-evans/create-or-update-comment@v5
         with:
           issue-number: ${{ github.event.pull_request.number }}
           body: |

HISTORY.md

@@ -4,6 +4,18 @@
 
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
 
+## 0.38.7
+
+Made a small tweak to DynamicPPL's compiler output to avoid potential undefined variables when resuming model functions midway through (e.g. with Libtask in Turing's SMC/PG samplers).
+
+## 0.38.6
+
+Renamed keyword argument `only_ddpl` to `only_dppl` for `Experimental.is_suitable_varinfo`.
+
+## 0.38.5
+
+Improve performance of VarNamedVector, mostly by changing how it handles contiguification.
+
 ## 0.38.4
 
 Improve performance of VarNamedVector. It should now be very nearly on par with Metadata for all models we've benchmarked on.

ext/DynamicPPLJETExt.jl

@@ -4,12 +4,12 @@ using DynamicPPL: DynamicPPL
 using JET: JET
 
 function DynamicPPL.Experimental.is_suitable_varinfo(
-    model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; only_ddpl::Bool=true
+    model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; only_dppl::Bool=true
 )
     f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(model, varinfo)
     # If specified, we only check errors originating somewhere in the DynamicPPL.jl.
     # This way we don't just fall back to untyped if the user's code is the issue.
-    result = if only_ddpl
+    result = if only_dppl
         JET.report_call(f, argtypes; target_modules=(JET.AnyFrameModule(DynamicPPL),))
     else
         JET.report_call(f, argtypes)
@@ -18,15 +18,15 @@ function DynamicPPL.Experimental.is_suitable_varinfo(
 end
 
 function DynamicPPL.Experimental._determine_varinfo_jet(
-    model::DynamicPPL.Model; only_ddpl::Bool=true
+    model::DynamicPPL.Model; only_dppl::Bool=true
 )
     # Generate a typed varinfo to test model type stability with
     varinfo = DynamicPPL.typed_varinfo(model)
 
     # Check type stability of evaluation (i.e. DefaultContext)
     model = DynamicPPL.setleafcontext(model, DynamicPPL.DefaultContext())
     eval_issuccess, eval_result = DynamicPPL.Experimental.is_suitable_varinfo(
-        model, varinfo; only_ddpl
+        model, varinfo; only_dppl
     )
     if !eval_issuccess
         @debug "Evaluation with typed varinfo failed with the following issues:"
@@ -36,7 +36,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     # Check type stability of initialisation (i.e. InitContext)
     model = DynamicPPL.setleafcontext(model, DynamicPPL.InitContext())
     init_issuccess, init_result = DynamicPPL.Experimental.is_suitable_varinfo(
-        model, varinfo; only_ddpl
+        model, varinfo; only_dppl
     )
     if !init_issuccess
         @debug "Initialisation with typed varinfo failed with the following issues:"

src/compiler.jl

@@ -461,9 +461,16 @@ function generate_tilde(left, right)
         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)(
+            # If `vn` is not in `argnames`, then it's definitely been conditioned on (if
+            # it's not in `argnames` and wasn't conditioned on, then `isassumption` would
+            # be true).
+            $left = if $(DynamicPPL.inargnames)($vn, __model__)
+                # This is a no-op and looks redundant, but defining the compiler output this
+                # way ensures that the variable `$left` is always defined. See
+                # https://github.com/TuringLang/DynamicPPL.jl/pull/1110.
+                $left
+            else
+                $(DynamicPPL.getconditioned_nested)(
                     __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
                 )
             end

src/experimental.jl

@@ -16,7 +16,7 @@ Check if the `model` supports evaluation using the provided `varinfo`.
 - `varinfo`: The varinfo to verify the support for.
 
 # Keyword Arguments
-- `only_ddpl`: If `true`, only consider error reports occuring in the tilde pipeline. Default: `true`.
+- `only_dppl`: If `true`, only consider error reports occuring in the tilde pipeline. Default: `true`.
 
 # Returns
 - `issuccess`: `true` if the model supports the varinfo, otherwise `false`.
@@ -28,7 +28,7 @@ function is_suitable_varinfo end
 function _determine_varinfo_jet end
 
 """
-    determine_suitable_varinfo(model; only_ddpl::Bool=true)
+    determine_suitable_varinfo(model; only_dppl::Bool=true)
 
 Return a suitable varinfo for the given `model`.
 
@@ -42,7 +42,7 @@ See also: [`DynamicPPL.Experimental.is_suitable_varinfo`](@ref).
 - `model`: The model for which to determine the varinfo.
 
 # Keyword Arguments
-- `only_ddpl`: If `true`, only consider error reports within DynamicPPL.jl.
+- `only_dppl`: If `true`, only consider error reports within DynamicPPL.jl.
 
 # Examples
 
@@ -83,10 +83,10 @@ julia> vi isa typeof(DynamicPPL.typed_varinfo(model_with_static_support()))
 true
 ```
 """
-function determine_suitable_varinfo(model::DynamicPPL.Model; only_ddpl::Bool=true)
+function determine_suitable_varinfo(model::DynamicPPL.Model; only_dppl::Bool=true)
     # If JET.jl has been loaded, and thus `determine_varinfo` has been defined, we use that.
     return if Base.get_extension(DynamicPPL, :DynamicPPLJETExt) !== nothing
-        _determine_varinfo_jet(model; only_ddpl)
+        _determine_varinfo_jet(model; only_dppl)
     else
         # Warn the user.
         @warn "JET.jl is not loaded. Assumes the model is compatible with typed varinfo."

src/varinfo.jl

@@ -1297,6 +1297,10 @@ function _link_metadata!!(
         metadata = setindex_internal!!(metadata, val_new, vn, transform_from_linked)
         set_transformed!(metadata, true, vn)
     end
+    # Linking can often change the sizes of variables, causing inactive elements. We don't
+    # want to keep them around, since typically linking is done once and then the VarInfo
+    # is evaluated multiple times. Hence we contiguify here.
+    metadata = contiguify!(metadata)
     return metadata, cumulative_logjac
 end
 
@@ -1465,6 +1469,10 @@ function _invlink_metadata!!(
         metadata = setindex_internal!!(metadata, tovec(new_val), vn, new_transform)
         set_transformed!(metadata, false, vn)
     end
+    # Linking can often change the sizes of variables, causing inactive elements. We don't
+    # want to keep them around, since typically linking is done once and then the VarInfo
+    # is evaluated multiple times. Hence we contiguify here.
+    metadata = contiguify!(metadata)
     return metadata, cumulative_inv_logjac
 end
 

src/varnamedvector.jl

@@ -341,10 +341,13 @@ function ==(vnv_left::VarNamedVector, vnv_right::VarNamedVector)
            vnv_left.num_inactive == vnv_right.num_inactive
 end
 
-function is_concretely_typed(vnv::VarNamedVector)
-    return isconcretetype(eltype(vnv.varnames)) &&
-           isconcretetype(eltype(vnv.vals)) &&
-           isconcretetype(eltype(vnv.transforms))
+function is_tightly_typed(vnv::VarNamedVector)
+    k = eltype(vnv.varnames)
+    v = eltype(vnv.vals)
+    t = eltype(vnv.transforms)
+    return (isconcretetype(k) || k === Union{}) &&
+           (isconcretetype(v) || v === Union{}) &&
+           (isconcretetype(t) || t === Union{})
 end
 
 getidx(vnv::VarNamedVector, vn::VarName) = vnv.varname_to_index[vn]
@@ -880,7 +883,16 @@ function loosen_types!!(
         return if vn_type == K && val_type == V && transform_type == T
             vnv
         elseif isempty(vnv)
-            VarNamedVector(vn_type[], val_type[], transform_type[])
+            VarNamedVector(
+                Dict{vn_type,Int}(),
+                Vector{vn_type}(),
+                UnitRange{Int}[],
+                Vector{val_type}(),
+                Vector{transform_type}(),
+                BitVector(),
+                Dict{Int,Int}();
+                check_consistency=false,
+            )
         else
             # TODO(mhauru) We allow a `vnv` to have any AbstractVector type as its vals, but
             # then here always revert to Vector.
@@ -944,7 +956,7 @@ julia> vnv_tight.transforms
 ```
 """
 function tighten_types!!(vnv::VarNamedVector)
-    return if is_concretely_typed(vnv)
+    return if is_tightly_typed(vnv)
         # There can not be anything to tighten, so short-circuit.
         vnv
     elseif isempty(vnv)
@@ -1020,6 +1032,7 @@ function insert_internal!!(
     end
     vnv = loosen_types!!(vnv, typeof(vn), eltype(val), typeof(transform))
     insert_internal!(vnv, val, vn, transform)
+    vnv = tighten_types!!(vnv)
     return vnv
 end
 
@@ -1029,6 +1042,7 @@ function update_internal!!(
     transform_resolved = transform === nothing ? gettransform(vnv, vn) : transform
     vnv = loosen_types!!(vnv, typeof(vn), eltype(val), typeof(transform_resolved))
     update_internal!(vnv, val, vn, transform)
+    vnv = tighten_types!!(vnv)
     return vnv
 end
 
@@ -1104,6 +1118,9 @@ care about them.
 
 This is in a sense the reverse operation of `vnv[:]`.
 
+The return value may share memory with the input `vnv`, and thus one can not be mutated
+safely without affecting the other.
+
 Unflatten recontiguifies the internal storage, getting rid of any inactive entries.
 
 # Examples
@@ -1125,15 +1142,20 @@ function unflatten(vnv::VarNamedVector, vals::AbstractVector)
             ),
         )
     end
-    new_ranges = deepcopy(vnv.ranges)
-    recontiguify_ranges!(new_ranges)
+    new_ranges = vnv.ranges
+    num_inactive = vnv.num_inactive
+    if has_inactive(vnv)
+        new_ranges = recontiguify_ranges!(new_ranges)
+        num_inactive = Dict{Int,Int}()
+    end
     return VarNamedVector(
         vnv.varname_to_index,
         vnv.varnames,
         new_ranges,
         vals,
         vnv.transforms,
-        vnv.is_unconstrained;
+        vnv.is_unconstrained,
+        num_inactive;
         check_consistency=false,
     )
 end
@@ -1428,6 +1450,9 @@ julia> vnv[@varname(x)]  # All the values are still there.
 ```
 """
 function contiguify!(vnv::VarNamedVector)
+    if !has_inactive(vnv)
+        return vnv
+    end
     # Extract the re-contiguified values.
     # NOTE: We need to do this before we update the ranges.
     old_vals = copy(vnv.vals)

test/Project.toml

@@ -4,6 +4,7 @@ AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
@@ -34,6 +35,7 @@ AbstractMCMC = "5"
 AbstractPPL = "0.13"
 Accessors = "0.1"
 Aqua = "0.8"
+BangBang = "0.4"
 Bijectors = "0.15.1"
 Combinatorics = "1"
 DifferentiationInterface = "0.6.41, 0.7"

test/accumulators.jl

@@ -117,7 +117,7 @@ using DynamicPPL:
             @test at_all64[:LogLikelihood] == ll_f64
 
             @test haskey(AccumulatorTuple(lp_f64), Val(:LogPrior))
-            @test ~haskey(AccumulatorTuple(lp_f64), Val(:LogLikelihood))
+            @test !haskey(AccumulatorTuple(lp_f64), Val(:LogLikelihood))
             @test length(AccumulatorTuple(lp_f64, ll_f64)) == 2
             @test keys(at_all64) == (:LogPrior, :LogLikelihood)
             @test collect(at_all64) == [lp_f64, ll_f64]

test/ext/DynamicPPLJETExt.jl

@@ -56,7 +56,7 @@
             DynamicPPL.NTVarInfo
         # Should fail if we're including errors in the model body.
         @test DynamicPPL.Experimental.determine_suitable_varinfo(
-            demo5(); only_ddpl=false
+            demo5(); only_dppl=false
         ) isa DynamicPPL.UntypedVarInfo
     end