Merge branch 'master' into sunxd/move_predict

Author: sunxd3

.github/workflows/CI.yml

@@ -19,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
 
@@ -60,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.5"
+version = "0.30.6"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/context_implementations.jl

@@ -77,49 +77,11 @@ function tilde_assume(
     return tilde_assume(rng, childcontext(context), args...)
 end
 
-function tilde_assume(context::PriorContext{<:NamedTuple}, right, vn, vi)
-    if haskey(context.vars, getsym(vn))
-        vi = setindex!!(vi, tovec(get(context.vars, vn)), vn)
-        settrans!!(vi, false, vn)
-    end
-    return tilde_assume(PriorContext(), right, vn, vi)
-end
-function tilde_assume(
-    rng::Random.AbstractRNG, context::PriorContext{<:NamedTuple}, sampler, right, vn, vi
-)
-    if haskey(context.vars, getsym(vn))
-        vi = setindex!!(vi, tovec(get(context.vars, vn)), vn)
-        settrans!!(vi, false, vn)
-    end
-    return tilde_assume(rng, PriorContext(), sampler, right, vn, vi)
-end
-
-function tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, vn, vi)
-    if haskey(context.vars, getsym(vn))
-        vi = setindex!!(vi, tovec(get(context.vars, vn)), vn)
-        settrans!!(vi, false, vn)
-    end
-    return tilde_assume(LikelihoodContext(), right, vn, vi)
-end
-function tilde_assume(
-    rng::Random.AbstractRNG,
-    context::LikelihoodContext{<:NamedTuple},
-    sampler,
-    right,
-    vn,
-    vi,
-)
-    if haskey(context.vars, getsym(vn))
-        vi = setindex!!(vi, tovec(get(context.vars, vn)), vn)
-        settrans!!(vi, false, vn)
-    end
-    return tilde_assume(rng, LikelihoodContext(), sampler, right, vn, vi)
-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)
@@ -257,7 +219,7 @@ function assume(
     else
         r = init(rng, dist, sampler)
         if istrans(vi)
-            f = to_linked_internal_transform(vi, dist)
+            f = to_linked_internal_transform(vi, vn, dist)
             push!!(vi, vn, f(r), dist, sampler)
             # By default `push!!` sets the transformed flag to `false`.
             settrans!!(vi, true, vn)
@@ -328,37 +290,6 @@ function dot_tilde_assume(
 end
 
 # `LikelihoodContext`
-function dot_tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, left, vn, vi)
-    return if haskey(context.vars, getsym(vn))
-        var = get(context.vars, vn)
-        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
-        set_val!(vi, _vns, _right, _left)
-        settrans!!.((vi,), false, _vns)
-        dot_tilde_assume(LikelihoodContext(), _right, _left, _vns, vi)
-    else
-        dot_tilde_assume(LikelihoodContext(), right, left, vn, vi)
-    end
-end
-function dot_tilde_assume(
-    rng::Random.AbstractRNG,
-    context::LikelihoodContext{<:NamedTuple},
-    sampler,
-    right,
-    left,
-    vn,
-    vi,
-)
-    return if haskey(context.vars, getsym(vn))
-        var = get(context.vars, vn)
-        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
-        set_val!(vi, _vns, _right, _left)
-        settrans!!.((vi,), false, _vns)
-        dot_tilde_assume(rng, LikelihoodContext(), sampler, _right, _left, _vns, vi)
-    else
-        dot_tilde_assume(rng, LikelihoodContext(), sampler, right, left, vn, vi)
-    end
-end
-
 function dot_tilde_assume(context::LikelihoodContext, right, left, vn, vi)
     return dot_assume(nodist(right), left, vn, vi)
 end
@@ -368,46 +299,16 @@ function dot_tilde_assume(
     return dot_assume(rng, sampler, nodist(right), vn, left, vi)
 end
 
-# `PriorContext`
-function dot_tilde_assume(context::PriorContext{<:NamedTuple}, right, left, vn, vi)
-    return if haskey(context.vars, getsym(vn))
-        var = get(context.vars, vn)
-        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
-        set_val!(vi, _vns, _right, _left)
-        settrans!!.((vi,), false, _vns)
-        dot_tilde_assume(PriorContext(), _right, _left, _vns, vi)
-    else
-        dot_tilde_assume(PriorContext(), right, left, vn, vi)
-    end
-end
-function dot_tilde_assume(
-    rng::Random.AbstractRNG,
-    context::PriorContext{<:NamedTuple},
-    sampler,
-    right,
-    left,
-    vn,
-    vi,
-)
-    return if haskey(context.vars, getsym(vn))
-        var = get(context.vars, vn)
-        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
-        set_val!(vi, _vns, _right, _left)
-        settrans!!.((vi,), false, _vns)
-        dot_tilde_assume(rng, PriorContext(), sampler, _right, _left, _vns, vi)
-    else
-        dot_tilde_assume(rng, PriorContext(), sampler, right, left, vn, vi)
-    end
-end
-
 # `PrefixContext`
 function dot_tilde_assume(context::PrefixContext, right, left, vn, vi)
-    return dot_tilde_assume(context.context, right, prefix.(Ref(context), vn), vi)
+    return dot_tilde_assume(context.context, right, left, prefix.(Ref(context), vn), vi)
 end
 
-function dot_tilde_assume(rng, context::PrefixContext, sampler, right, left, vn, vi)
+function dot_tilde_assume(
+    rng::Random.AbstractRNG, context::PrefixContext, sampler, right, left, vn, vi
+)
     return dot_tilde_assume(
-        rng, context.context, sampler, right, prefix.(Ref(context), vn), vi
+        rng, context.context, sampler, right, left, prefix.(Ref(context), vn), vi
     )
 end
 
@@ -500,7 +401,7 @@ end
 # HACK: These methods are only used in the `get_and_set_val!` methods below.
 # FIXME: Remove these.
 function _link_broadcast_new(vi, vn, dist, r)
-    b = to_linked_internal_transform(vi, dist)
+    b = to_linked_internal_transform(vi, vn, dist)
     return b(r)
 end
 
@@ -591,7 +492,7 @@ function get_and_set_val!(
             push!!.((vi,), vns, _link_broadcast_new.((vi,), vns, dists, r), dists, (spl,))
             # NOTE: Need to add the correction.
             # FIXME: This is not great.
-            acclogp_assume!!(vi, sum(logabsdetjac.(link_transform.(dists), r)))
+            acclogp!!(vi, sum(logabsdetjac.(link_transform.(dists), r)))
             # `push!!` sets the trans-flag to `false` by default.
             settrans!!.((vi,), true, vns)
         else

src/contexts.jl

@@ -53,7 +53,7 @@ DefaultContext()
 julia> ctx_prior = DynamicPPL.setchildcontext(ctx, PriorContext()); # only compute the logprior
 
 julia> DynamicPPL.childcontext(ctx_prior)
-PriorContext{Nothing}(nothing)
+PriorContext()
 ```
 """
 setchildcontext
@@ -97,7 +97,7 @@ ParentContext(ParentContext(DefaultContext()))
 
 julia> # Replace the leaf context with another leaf.
        leafcontext(setleafcontext(ctx, PriorContext()))
-PriorContext{Nothing}(nothing)
+PriorContext()
 
 julia> # Append another parent context.
        setleafcontext(ctx, ParentContext(DefaultContext()))
@@ -195,32 +195,19 @@ struct DefaultContext <: AbstractContext end
 NodeTrait(context::DefaultContext) = IsLeaf()
 
 """
-    struct PriorContext{Tvars} <: AbstractContext
-        vars::Tvars
-    end
+    PriorContext <: AbstractContext
 
-The `PriorContext` enables the computation of the log prior of the parameters `vars` when
-running the model.
+A leaf context resulting in the exclusion of likelihood terms when running the model.
 """
-struct PriorContext{Tvars} <: AbstractContext
-    vars::Tvars
-end
-PriorContext() = PriorContext(nothing)
+struct PriorContext <: AbstractContext end
 NodeTrait(context::PriorContext) = IsLeaf()
 
 """
-    struct LikelihoodContext{Tvars} <: AbstractContext
-        vars::Tvars
-    end
+    LikelihoodContext <: AbstractContext
 
-The `LikelihoodContext` enables the computation of the log likelihood of the parameters when
-running the model. `vars` can be used to evaluate the log likelihood for specific values
-of the model's parameters. If `vars` is `nothing`, the parameter values inside the `VarInfo` will be used by default.
+A leaf context resulting in the exclusion of prior terms when running the model.
 """
-struct LikelihoodContext{Tvars} <: AbstractContext
-    vars::Tvars
-end
-LikelihoodContext() = LikelihoodContext(nothing)
+struct LikelihoodContext <: AbstractContext end
 NodeTrait(context::LikelihoodContext) = IsLeaf()
 
 """

src/debug_utils.jl

@@ -332,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
@@ -425,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(

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})