Merge branch 'main' into breaking

Author: penelopeysm

.github/workflows/IntegrationTest.yml

@@ -43,8 +43,6 @@ jobs:
 
       - name: Load this and run the downstream tests
         shell: julia --color=yes --project=downstream {0}
-        # Don't test Turing.jl on 1.12.0, it's broken because of Libtask
-        if: ${{ !(steps.julia-version.outputs.julia == '1.12.0' && matrix.package.repo == 'Turing.jl') }}
         run: |
           using Pkg
           try

HISTORY.md

@@ -2,6 +2,21 @@
 
 ## 0.39.0
 
+Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
+
+## 0.38.3
+
+Add an implementation of `returned(::Model, ::AbstractDict{<:VarName})`.
+Please note we generally recommend using Dict, as NamedTuples cannot correctly represent variables with indices / fields on the left-hand side of tildes, like `x[1]` or `x.a`.
+
+The generic method `returned(::Model, values, keys)` is deprecated and will be removed in the next minor version.
+
+## 0.38.2
+
+Added a compatibility entry for JET@0.11.
+
+> > > > > > > main
+
 ## 0.38.1
 
 Added `from_linked_vec_transform` and `from_vec_transform` methods for `ProductNamedTupleDistribution`.

Project.toml

@@ -62,7 +62,7 @@ DocStringExtensions = "0.9"
 EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10.12, 1"
 InteractiveUtils = "1"
-JET = "0.9, 0.10"
+JET = "0.9, 0.10, 0.11"
 KernelAbstractions = "0.9.33"
 LinearAlgebra = "1.6"
 LogDensityProblems = "2"

docs/Project.toml

@@ -22,7 +22,7 @@ DocumenterMermaid = "0.1, 0.2"
 DynamicPPL = "0.38"
 FillArrays = "0.13, 1"
 ForwardDiff = "0.10, 1"
-JET = "0.9, 0.10"
+JET = "0.9, 0.10, 0.11"
 LogDensityProblems = "2"
 MarginalLogDensities = "0.4"
 MCMCChains = "5, 6, 7"

docs/src/api.md

@@ -176,11 +176,11 @@ It is possible to manually increase (or decrease) the accumulated log likelihood
 @addlogprob!
 ```
 
-Return values of the model function can be obtained with [`returned(model, sample)`](@ref), where `sample` is either a `MCMCChains.Chains` object (which represents a collection of samples) or a single sample represented as a `NamedTuple`.
+Return values of the model function can be obtained with [`returned(model, sample)`](@ref), where `sample` is either a `MCMCChains.Chains` object (which represents a collection of samples), or a single sample represented as a `NamedTuple` or a dictionary of VarNames.
 
 ```@docs
 returned(::DynamicPPL.Model, ::MCMCChains.Chains)
-returned(::DynamicPPL.Model, ::NamedTuple)
+returned(::DynamicPPL.Model, ::Union{NamedTuple,AbstractDict{<:VarName}})
 ```
 
 For a chain of samples, one can compute the pointwise log-likelihoods of each observed random variable with [`pointwise_loglikelihoods`](@ref). Similarly, the log-densities of the priors using

src/model.jl

@@ -1103,44 +1103,36 @@ function predict end
 
 """
     returned(model::Model, parameters::NamedTuple)
-    returned(model::Model, values, keys)
-    returned(model::Model, values, keys)
+    returned(model::Model, parameters::AbstractDict{<:VarName})
 
 Execute `model` with variables `keys` set to `values` and return the values returned by the `model`.
 
-If a `NamedTuple` is given, `keys=keys(parameters)` and `values=values(parameters)`.
-
 # Example
 ```jldoctest
 julia> using DynamicPPL, Distributions
 
-julia> @model function demo(xs)
-           s ~ InverseGamma(2, 3)
-           m_shifted ~ Normal(10, √s)
-           m = m_shifted - 10
-           for i in eachindex(xs)
-               xs[i] ~ Normal(m, √s)
-           end
-           return (m, )
+julia> @model function demo()
+           m ~ Normal()
+           return (mp1 = m + 1,)
        end
 demo (generic function with 2 methods)
 
-julia> model = demo(randn(10));
-
-julia> parameters = (; s = 1.0, m_shifted=10.0);
+julia> model = demo();
 
-julia> returned(model, parameters)
-(0.0,)
+julia> returned(model, (; m = 1.0))
+(mp1 = 2.0,)
 
-julia> returned(model, values(parameters), keys(parameters))
-(0.0,)
+julia> returned(model, Dict{VarName,Float64}(@varname(m) => 2.0))
+(mp1 = 3.0,)
 ```
 """
-function returned(model::Model, parameters::NamedTuple)
-    fixed_model = fix(model, parameters)
-    return fixed_model()
-end
-
-function returned(model::Model, values, keys)
-    return returned(model, NamedTuple{keys}(values))
+function returned(model::Model, parameters::Union{NamedTuple,AbstractDict{<:VarName}})
+    vi = DynamicPPL.setaccs!!(VarInfo(), ())
+    # Note: we can't use `fix(model, parameters)` because
+    # https://github.com/TuringLang/DynamicPPL.jl/issues/1097
+    # Use `nothing` as the fallback to ensure that any missing parameters cause an error
+    ctx = InitContext(Random.default_rng(), InitFromParams(parameters, nothing))
+    new_model = setleafcontext(model, ctx)
+    # We can't use new_model() because that overwrites it with an InitContext of its own.
+    return first(evaluate!!(new_model, vi))
 end

test/Project.toml

@@ -41,7 +41,7 @@ Distributions = "0.25"
 DistributionsAD = "0.6.3"
 Documenter = "1"
 ForwardDiff = "0.10.12, 1"
-JET = "0.9, 0.10"
+JET = "0.9, 0.10, 0.11"
 LogDensityProblems = "2"
 MCMCChains = "7.2.1"
 MacroTools = "0.5.6"

test/model.jl

@@ -321,6 +321,38 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
         end
     end
 
+    @testset "returned() on NamedTuple / Dict" begin
+        @model function demo_returned()
+            a ~ Normal()
+            b ~ Normal()
+            return (asq=a^2, bsq=b^2)
+        end
+        model = demo_returned()
+
+        @testset "NamedTuple" begin
+            params = (a=1.0, b=2.0)
+            results = returned(model, params)
+            @test results.asq == params.a^2
+            @test results.bsq == params.b^2
+            # `returned` should error when not all parameters are provided
+            @test_throws ErrorException returned(model, (; a=1.0))
+            @test_throws ErrorException returned(model, (a=1.0, b=missing))
+        end
+        @testset "Dict" begin
+            params = Dict{VarName,Float64}(@varname(a) => 1.0, @varname(b) => 2.0)
+            results = returned(model, params)
+            @test results.asq == params[@varname(a)]^2
+            @test results.bsq == params[@varname(b)]^2
+            # `returned` should error when not all parameters are provided
+            @test_throws ErrorException returned(
+                model, Dict{VarName,Float64}(@varname(a) => 1.0)
+            )
+            @test_throws ErrorException returned(
+                model, Dict{VarName,Any}(@varname(a) => 1.0, @varname(b) => missing)
+            )
+        end
+    end
+
     @testset "returned() on `LKJCholesky`" begin
         n = 10
         d = 2