Merge branch 'master' into torfjelde/fix-fixes

Author: penelopeysm

.github/workflows/CI.yml

@@ -4,9 +4,11 @@ on:
   push:
     branches:
       - master
+      - backport-*
   pull_request:
     branches:
       - master
+      - backport-*
   merge_group:
     types: [checks_requested]
 
@@ -17,39 +19,53 @@ permissions:
 
 jobs:
   test:
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ matrix.runner.os }}
     strategy:
       matrix:
-        version:
-          - 'min' # minimum supported version
-          - '1' # current stable version
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-        num_threads:
-          - 1
-          - 2
-        include:
+        runner:
+          # Current stable version
+          - version: '1'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 2
+          # Minimum supported version
+          - version: 'min'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 2
+          # Single-threaded
+          - version: '1'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 1
+          # Minimum supported version, single-threaded
+          - version: 'min'
+            os: ubuntu-latest
+            arch: x64
+            num_threads: 1
+          # x86
           - version: '1'
             os: ubuntu-latest
             arch: x86
             num_threads: 2
+          # Windows
           - version: '1'
             os: windows-latest
             arch: x64
             num_threads: 2
+          # macOS
           - version: '1'
-            os: macOS-latest
-            arch: x64
+            os: macos-latest
+            arch: aarch64
             num_threads: 2
+
     steps:
       - uses: actions/checkout@v4
 
       - uses: julia-actions/setup-julia@v2
         with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
+          version: ${{ matrix.runner.version }}
+          arch: ${{ matrix.runner.arch }}
 
       - uses: julia-actions/cache@v2
 
@@ -58,7 +74,7 @@ jobs:
       - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: All
-          JULIA_NUM_THREADS: ${{ matrix.num_threads }}
+          JULIA_NUM_THREADS: ${{ matrix.runner.num_threads }}
 
       - uses: julia-actions/julia-processcoverage@v1
 

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.30.2"
+version = "0.30.6"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -14,6 +14,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -46,7 +47,7 @@ AbstractMCMC = "5"
 AbstractPPL = "0.8.4, 0.9"
 Accessors = "0.1"
 BangBang = "0.4.1"
-Bijectors = "0.13.18"
+Bijectors = "0.13.18, 0.14"
 ChainRulesCore = "1"
 Compat = "4"
 ConstructionBase = "1.5.4"
@@ -65,7 +66,7 @@ Requires = "1"
 ReverseDiff = "1"
 Test = "1.6"
 ZygoteRules = "0.2"
-julia = "~1.6.6, 1.7.3"
+julia = "1.10"
 
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/api.md

@@ -212,6 +212,19 @@ And some which might be useful to determine certain properties of the model base
 DynamicPPL.has_static_constraints
 ```
 
+For determining whether one might have type instabilities in the model, the following can be useful
+
+```@docs
+DynamicPPL.DebugUtils.model_warntype
+DynamicPPL.DebugUtils.model_typed
+```
+
+Interally, the type-checking methods make use of the following method for construction of the call with the argument types:
+
+```@docs
+DynamicPPL.DebugUtils.gen_evaluator_call_with_types
+```
+
 ## Advanced
 
 ### Variable names

ext/DynamicPPLMCMCChainsExt.jl

@@ -108,86 +108,14 @@ function DynamicPPL.generated_quantities(
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     return map(iters) do (sample_idx, chain_idx)
-        if DynamicPPL.supports_varname_indexing(chain)
-            varname_pairs = _varname_pairs_with_varname_indexing(
-                chain, varinfo, sample_idx, chain_idx
-            )
-        else
-            varname_pairs = _varname_pairs_without_varname_indexing(
-                chain, varinfo, sample_idx, chain_idx
-            )
-        end
-        fixed_model = DynamicPPL.fix(model, Dict(varname_pairs))
-        return fixed_model()
+        # TODO: Use `fix` once we've addressed https://github.com/TuringLang/DynamicPPL.jl/issues/702.
+        # Update the varinfo with the current sample and make variables not present in `chain`
+        # to be sampled.
+        DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
+        # NOTE: Some of the varialbes can be a view into the `varinfo`, so we need to
+        # `deepcopy` the `varinfo` before passing it to the `model`.
+        model(deepcopy(varinfo))
     end
 end
 
-"""
-    _varname_pairs_with_varname_indexing(
-        chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
-    )
-
-Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values
-from the chain.
-
-This implementation assumes `chain` can be indexed using variable names, and is the
-preffered implementation.
-"""
-function _varname_pairs_with_varname_indexing(
-    chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
-)
-    vns = DynamicPPL.varnames(chain)
-    vn_parents = Iterators.map(vns) do vn
-        # The call nested_setindex_maybe! is used to handle cases where vn is not
-        # the variable name used in the model, but rather subsumed by one. Except
-        # for the subsumption part, this could be
-        # vn => getindex_varname(chain, sample_idx, vn, chain_idx)
-        # TODO(mhauru) This call to nested_setindex_maybe! is unintuitive.
-        DynamicPPL.nested_setindex_maybe!(
-            varinfo, DynamicPPL.getindex_varname(chain, sample_idx, vn, chain_idx), vn
-        )
-    end
-    varname_pairs = Iterators.map(Iterators.filter(!isnothing, vn_parents)) do vn_parent
-        vn_parent => varinfo[vn_parent]
-    end
-    return varname_pairs
-end
-
-"""
-Check which keys in `key_strings` are subsumed by `vn_string` and return the their values.
-
-The subsumption check is done with `DynamicPPL.subsumes_string`, which is quite weak, and
-won't catch all cases. We should get rid of this if we can.
-"""
-# TODO(mhauru) See docstring above.
-function _vcat_subsumed_values(vn_string, values, key_strings)
-    indices = findall(Base.Fix1(DynamicPPL.subsumes_string, vn_string), key_strings)
-    return !isempty(indices) ? reduce(vcat, values[indices]) : nothing
-end
-
-"""
-    _varname_pairs_without_varname_indexing(
-        chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
-    )
-
-Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values
-from the chain.
-
-This implementation does not assume that `chain` can be indexed using variable names. It is
-thus not guaranteed to work in cases where the variable names have complex subsumption
-patterns, such as if the model has a variable `x` but the chain stores `x.a[1]`.
-"""
-function _varname_pairs_without_varname_indexing(
-    chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx
-)
-    values = chain.value[sample_idx, :, chain_idx]
-    keys = Base.keys(chain)
-    keys_strings = map(string, keys)
-    varname_pairs = [
-        vn => _vcat_subsumed_values(string(vn), values, keys_strings) for
-        vn in Base.keys(varinfo)
-    ]
-    return varname_pairs
-end
-
 end

src/context_implementations.jl

@@ -78,10 +78,10 @@ function tilde_assume(
 end
 
 function tilde_assume(::LikelihoodContext, right, vn, vi)
-    return assume(NoDist(right), vn, vi)
+    return assume(nodist(right), vn, vi)
 end
 function tilde_assume(rng::Random.AbstractRNG, ::LikelihoodContext, sampler, right, vn, vi)
-    return assume(rng, sampler, NoDist(right), vn, vi)
+    return assume(rng, sampler, nodist(right), vn, vi)
 end
 
 function tilde_assume(context::PrefixContext, right, vn, vi)

src/debug_utils.jl

@@ -5,6 +5,7 @@ using ..DynamicPPL: broadcast_safe, AbstractContext, childcontext
 
 using Random: Random
 using Accessors: Accessors
+using InteractiveUtils: InteractiveUtils
 
 using DocStringExtensions
 using Distributions
@@ -331,7 +332,9 @@ function DynamicPPL.tilde_assume(context::DebugContext, right, vn, vi)
     record_post_tilde_assume!(context, vn, right, value, logp, vi)
     return value, logp, vi
 end
-function DynamicPPL.tilde_assume(rng, context::DebugContext, sampler, right, vn, vi)
+function DynamicPPL.tilde_assume(
+    rng::Random.AbstractRNG, context::DebugContext, sampler, right, vn, vi
+)
     record_pre_tilde_assume!(context, vn, right, vi)
     value, logp, vi = DynamicPPL.tilde_assume(
         rng, childcontext(context), sampler, right, vn, vi
@@ -424,7 +427,7 @@ function DynamicPPL.dot_tilde_assume(context::DebugContext, right, left, vn, vi)
 end
 
 function DynamicPPL.dot_tilde_assume(
-    rng, context::DebugContext, sampler, right, left, vn, vi
+    rng::Random.AbstractRNG, context::DebugContext, sampler, right, left, vn, vi
 )
     record_pre_dot_tilde_assume!(context, vn, left, right, vi)
     value, logp, vi = DynamicPPL.dot_tilde_assume(
@@ -678,4 +681,83 @@ function has_static_constraints(
     return all_the_same(transforms)
 end
 
+"""
+    gen_evaluator_call_with_types(model[, varinfo, context])
+
+Generate the evaluator call and the types of the arguments.
+
+# Arguments
+- `model::Model`: The model whose evaluator is of interest.
+- `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
+- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
+
+# Returns
+A 2-tuple with the following elements:
+- `f`: This is either `model.f` or `Core.kwcall`, depending on whether
+    the model has keyword arguments.
+- `argtypes::Type{<:Tuple}`: The types of the arguments for the evaluator.
+"""
+function gen_evaluator_call_with_types(
+    model::Model,
+    varinfo::AbstractVarInfo=VarInfo(model),
+    context::AbstractContext=DefaultContext(),
+)
+    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
+    return if isempty(kwargs)
+        (model.f, Base.typesof(args...))
+    else
+        (Core.kwcall, Tuple{typeof(kwargs),Core.Typeof(model.f),map(Core.Typeof, args)...})
+    end
+end
+
+"""
+    model_warntype(model[, varinfo, context]; optimize=true)
+
+Check the type stability of the model's evaluator, warning about any potential issues.
+
+This simply calls `@code_warntype` on the model's evaluator, filling in internal arguments where needed.
+
+# Arguments
+- `model::Model`: The model to check.
+- `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
+- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
+
+# Keyword Arguments
+- `optimize::Bool`: Whether to generate optimized code. Default: `false`.
+"""
+function model_warntype(
+    model::Model,
+    varinfo::AbstractVarInfo=VarInfo(model),
+    context::AbstractContext=DefaultContext();
+    optimize::Bool=false,
+)
+    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo, context)
+    return InteractiveUtils.code_warntype(ftype, argtypes; optimize=optimize)
+end
+
+"""
+    model_typed(model[, varinfo, context]; optimize=true)
+
+Return the type inference for the model's evaluator.
+
+This simply calls `@code_typed` on the model's evaluator, filling in internal arguments where needed.
+
+# Arguments
+- `model::Model`: The model to check.
+- `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
+- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
+
+# Keyword Arguments
+- `optimize::Bool`: Whether to generate optimized code. Default: `true`.
+"""
+function model_typed(
+    model::Model,
+    varinfo::AbstractVarInfo=VarInfo(model),
+    context::AbstractContext=DefaultContext();
+    optimize::Bool=true,
+)
+    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo, context)
+    return only(InteractiveUtils.code_typed(ftype, argtypes; optimize=optimize))
+end
+
 end

src/distribution_wrappers.jl

@@ -42,6 +42,14 @@ Base.length(dist::NoDist) = Base.length(dist.dist)
 Base.size(dist::NoDist) = Base.size(dist.dist)
 
 Distributions.rand(rng::Random.AbstractRNG, d::NoDist) = rand(rng, d.dist)
+# NOTE(torfjelde): Need this to avoid stack overflow.
+function Distributions.rand!(
+    rng::Random.AbstractRNG,
+    d::NoDist{Distributions.ArrayLikeVariate{N}},
+    x::AbstractArray{<:Real,N},
+) where {N}
+    return Distributions.rand!(rng, d.dist, x)
+end
 Distributions.logpdf(d::NoDist{<:Univariate}, ::Real) = 0
 Distributions.logpdf(d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}) = 0
 function Distributions.logpdf(d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real})