From 2d8df0b78e793447fb3c521090db5f041b697956 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 20 Oct 2026 23:13:38 +0100 Subject: [PATCH 1/7] threadsafe: reject tilde assignments that replace containers Assisted-by: Codex --- HISTORY.md | 2 + docs/src/accs/threadsafe.md | 16 +++ src/compiler.jl | 14 +-- src/threadsafe.jl | 64 ++++++++++++ test/threadsafe.jl | 190 +++++++++++++++++++++++++++++++++++- 5 files changed, 274 insertions(+), 12 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index a6debdd0db..4292074f58 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,5 +1,7 @@ # 0.43.0 (unreleased) +Thread-safe evaluation now throws an `ArgumentError` when a tilde assignment replaces a mutable container or changes an immutable container's type; allocate mutable storage with `::Type{T}=Float64` and `Vector{T}(undef, n)`. + `VarNamedTuple` membership resolves `begin` and `end` against the stored array. See [#1490](https://github.com/TuringLang/DynamicPPL.jl/pull/1490). `ComponentVector` properties, including nested fields and slices, use consistent indices for membership, retrieval, and updates. See [#1491](https://github.com/TuringLang/DynamicPPL.jl/pull/1491). diff --git a/docs/src/accs/threadsafe.md b/docs/src/accs/threadsafe.md index 3c60b25f27..41a20eb032 100644 --- a/docs/src/accs/threadsafe.md +++ b/docs/src/accs/threadsafe.md @@ -75,3 +75,19 @@ Parameter values are also accumulator outputs. `VectorValueAccumulator` and `RawValueAccumulator` use task-local storage and combine their results after evaluation. This does not make arbitrary mutations in the model body or shared mutable initialisation strategies thread-safe. + +## Containers under thread-safe evaluation + +Under thread-safe evaluation, a tilde assignment throws an `ArgumentError` if it +replaces a mutable container on its path or changes an immutable container's type, +wherever the assignment occurs. Slice reads are checked by type because they return +copies. Allocate with the parameter element type via `::Type{T}=Float64`: use +`Vector{T}(undef, n)` for scalar tildes or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` +for vector-valued tildes. Assignments into data arguments that cannot be promoted also +throw, including `Vector{Missing}` and integer data imputed by a continuous +distribution; `Union{Missing,Float32}` data can still throw during initialisation from +the prior. Use `Union{Missing,Float64}` data for continuous imputation. The check does +not prevent lost writes from rebinding shared immutable values, including `NamedTuple`s +and immutable wrappers around mutable storage such as `OffsetArray`s of `Vector{Any}`, +or from concurrent writes through overlapping slices (`x[1:N][i] ~`). Serial evaluation +with thread-safe mode disabled retains its usual widening behaviour. diff --git a/src/compiler.jl b/src/compiler.jl index efd9a622fc..f3547e350e 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -431,10 +431,8 @@ function assign_or_set!!(lhs::Expr, rhs, vn) left_top_sym = get_top_level_symbol(lhs) return drop_escape( :( - $left_top_sym = $(Accessors.set)( - $left_top_sym, - $(AbstractPPL.with_mutation)($(AbstractPPL.getoptic)($vn)), - $rhs, + $left_top_sym = $(DynamicPPL.set_tilde_value!!)( + __varinfo__, $left_top_sym, $rhs, $vn ) ), ) @@ -466,7 +464,7 @@ function generate_tilde(left, right) # tilde_observe!! for more details. if $(DynamicPPL.isfixed(left, vn)) && !($dist isa $(DynamicPPL.Submodel)) # $left may not be a simple varname, it might be x.a or x[1], in which case we - # need to use Accessors.set to safely set it. + # need to use the tilde setter to safely set it. $(assign_or_set!!( left, :($(DynamicPPL.getfixed_nested)( @@ -561,10 +559,8 @@ function generate_tilde_assume(left, right, vn) expr = if left isa Expr # as opposed to Symbol left_top_sym = get_top_level_symbol(left) :( - $left_top_sym = $(Accessors.set)( - $left_top_sym, - $(AbstractPPL.with_mutation)($(AbstractPPL.getoptic)($vn)), - $value, + $left_top_sym = $(DynamicPPL.set_tilde_value!!)( + __varinfo__, $left_top_sym, $value, $vn ) ) else diff --git a/src/threadsafe.jl b/src/threadsafe.jl index 45e9712b29..d34751ffe4 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -44,6 +44,70 @@ function _threadsafe_varinfo(vi::AbstractVarInfo, ::Type{L}) where {L<:Accumulat end ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi +function set_tilde_value!!(::AbstractVarInfo, object, value, vn::VarName) + return Accessors.set(object, AbstractPPL.with_mutation(AbstractPPL.getoptic(vn)), value) +end + +@noinline function _tilde_container_error(old, new, vn) + replacement = if typeof(old) === typeof(new) + "$(typeof(old)) was replaced by a copy of the same type. " + else + "$(typeof(old)) -> $(typeof(new)). " + end + throw( + ArgumentError( + "Thread-safe evaluation cannot replace the container of tilde variable " * + "$(vn): " * + replacement * + "Mutable containers must retain their identity; immutable containers must " * + "retain their type. Allocate with the parameter element type (for example, " * + "`::Type{T}=Float64` with `Vector{T}(undef, n)` for scalar tildes, " * + "or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes).", + ), + ) +end + +@inline function _capture_tilde_path(old, optic) + tail = AbstractPPL.otail(optic) + tail isa AbstractPPL.Iden && return ((old, false),) + sliced = _tilde_reads_slice(optic, old) + # The final slice copy is written back into its checked parent; reading it here + # would only allocate another copy. + if AbstractPPL.otail(tail) isa AbstractPPL.Iden && sliced + return ((old, sliced),) + end + child = AbstractPPL.ohead(optic)(old) + return ((old, sliced), _capture_tilde_path(child, tail)...) +end + +function _tilde_reads_slice(optic::AbstractPPL.Index, object::AbstractArray) + indices = AbstractPPL.concretize(AbstractPPL.ohead(optic), object) + return any( + i -> i isa Union{AbstractArray,Colon}, (indices.ix..., values(indices.kw)...) + ) +end +_tilde_reads_slice(optic, object) = false + +@inline function _check_tilde_path(old_path, new, optic, vn, sliced=false) + old, child_sliced = first(old_path) + valid = ismutable(old) && !sliced ? old === new : typeof(old) === typeof(new) + valid || _tilde_container_error(old, new, vn) + tail = AbstractPPL.otail(optic) + tail isa AbstractPPL.Iden && return nothing + AbstractPPL.otail(tail) isa AbstractPPL.Iden && child_sliced && return nothing + return _check_tilde_path( + Base.tail(old_path), AbstractPPL.ohead(optic)(new), tail, vn, child_sliced + ) +end + +function set_tilde_value!!(::ThreadSafeVarInfo, object, value, vn::VarName) + optic = AbstractPPL.getoptic(vn) + old_path = _capture_tilde_path(object, optic) + new = Accessors.set(object, AbstractPPL.with_mutation(optic), value) + _check_tilde_path(old_path, new, optic, vn) + return new +end + function _get_task_accs_cache(vi::ThreadSafeVarInfo{V,L}) where {V,L} cache = @atomic :acquire vi.task_accs_cache # A serialized cache may have been sized for a process with fewer threads. diff --git a/test/threadsafe.jl b/test/threadsafe.jl index 21e11d8da9..ac39fa4b08 100644 --- a/test/threadsafe.jl +++ b/test/threadsafe.jl @@ -9,7 +9,8 @@ using DynamicPPL using Test using ForwardDiff using Random: Xoshiro -using LogDensityProblems: logdensity +using ADTypes: AutoForwardDiff +using LogDensityProblems: logdensity, logdensity_and_gradient @model function gdemo_d() s ~ InverseGamma(2, 3) @@ -146,8 +147,8 @@ end end @testset "parameter containers preserve numeric types" begin - @model function indexed_parameter() - x = zeros(1) + @model function indexed_parameter(::Type{T}=Float64) where {T} + x = zeros(T, 1) return x[1] ~ Normal() end model = setthreadsafe(indexed_parameter(), true) @@ -420,6 +421,189 @@ end @test !check_model(model) end + @testset "threadsafe assignments reject container replacement" begin + @model function shared() + x = zeros(8) + Threads.@threads for i in eachindex(x) + x[i] ~ Normal() + end + return 0.5 ~ Normal(sum(x), 1) + end + @model function serial() + x = zeros(8) + for i in eachindex(x) + x[i] ~ Normal() + end + return 0.5 ~ Normal(sum(x), 1) + end + @model function same_type_copy() + x = Vector{Real}[zeros(2) for _ in 1:8] + Threads.@threads for i in eachindex(x) + x[i] ~ MvNormal(zeros(2), ones(2)) + end + return 0.5 ~ Normal(sum(sum, x), 1) + end + @model function dictionary() + key = [1] + x = Dict(key => Vector{Real}[zeros(1)]) + return x[key][1] ~ MvNormal(ones(1), ones(1)) + end + @test_throws ArgumentError init!!( + Xoshiro(1), + setthreadsafe(dictionary(), true), + VarInfo(LogPriorAccumulator()), + InitFromPrior(), + UnlinkAll(), + ) + @model function nested() + x = Any[zeros(8)] + Threads.@threads for j in eachindex(x[1]) + x[1][j] ~ Normal() + end + return 0.5 ~ Normal(sum(x[1]), 1) + end + mutable struct Box + v + end + @model function boxed() + x = Box(zeros(1)) + Threads.@threads for j in eachindex(x.v) + x.v[j] ~ Normal() + end + return 0.5 ~ Normal(sum(x.v), 1) + end + @model function typed(::Type{T}=Float64) where {T} + x = Vector{T}(undef, 8) + Threads.@threads for i in eachindex(x) + x[i] ~ Normal() + end + return 0.5 ~ Normal(sum(x), 1) + end + @model function local_values() + Threads.@threads for i in 1:8 + s = (a=0.0,) + s.a ~ Normal() + end + end + + function nested_exception(err) + while err isa CompositeException || err isa TaskFailedException + err = err isa CompositeException ? first(err.exceptions) : err.task.result + end + return err + end + function check_type_change(model) + model = setthreadsafe(model, true) + err = try + vi = VarInfo(Xoshiro(1), model) + ldf = LogDensityFunction( + model, getlogjoint_internal, vi; adtype=AutoForwardDiff() + ) + logdensity_and_gradient(ldf, vi[:]) + nothing + catch err + nested_exception(err) + end + @test err isa ArgumentError + message = sprint(showerror, err) + @test occursin(r"tilde variable x(\.v)?\[", message) + @test occursin("Vector{T}(undef, n)", message) + end + for model in (shared(), same_type_copy(), nested(), boxed()) + check_type_change(model) + end + + serial_model = serial() + serial_ldf = LogDensityFunction(serial_model; adtype=AutoForwardDiff()) + @test last(logdensity_and_gradient(serial_ldf, zeros(8))) ≈ fill(0.5, 8) + @test_throws ArgumentError logdensity_and_gradient( + LogDensityFunction(setthreadsafe(serial_model, true); adtype=AutoForwardDiff()), + zeros(8), + ) + + typed_model = setthreadsafe(typed(), true) + typed_ldf = LogDensityFunction(typed_model; adtype=AutoForwardDiff()) + @test last(logdensity_and_gradient(typed_ldf, zeros(8))) ≈ fill(0.5, 8) + + @model function sliced(indices, ::Type{T}=Float64) where {T} + x = Matrix{T}(undef, 2, 4) + Threads.@threads for j in axes(x, 2) + x[indices, j] .~ Normal() + end + return 0.5 ~ Normal(sum(x), 1) + end + @model function columns(::Type{T}=Float64) where {T} + x = Matrix{T}(undef, 2, 4) + Threads.@threads for j in axes(x, 2) + x[:, j] ~ MvNormal(zeros(2), ones(2)) + end + return 0.5 ~ Normal(sum(x), 1) + end + for model in (sliced(:), sliced(1:2), sliced([1, 2]), columns()) + model = setthreadsafe(model, true) + @test logdensity(LogDensityFunction(model), zeros(8)) ≈ + 8 * logpdf(Normal(), 0.0) + logpdf(Normal(), 0.5) + @test last( + logdensity_and_gradient( + LogDensityFunction(model; adtype=AutoForwardDiff()), zeros(8) + ), + ) ≈ fill(0.5, 8) + end + + @model function sliced_observation(y) + mu ~ Normal() + Threads.@threads for j in axes(y, 2) + y[:, j] .~ Normal(mu, 1) + end + end + y = reshape(collect(1.0:4.0), 2, 2) + observed_model = setthreadsafe(sliced_observation(y), true) + @test logdensity(LogDensityFunction(observed_model), [0.0]) ≈ + logpdf(Normal(), 0.0) + sum(logpdf.(Normal(), y)) + @test last( + logdensity_and_gradient( + LogDensityFunction(observed_model; adtype=AutoForwardDiff()), [0.0] + ), + ) ≈ [sum(y)] + + local_model = setthreadsafe(local_values(), true) + local_ldf = LogDensityFunction(local_model) + @test logdensity(local_ldf, zeros(1)) ≈ 8 * logpdf(Normal(), 0.0) + # Accepted limitation: task-local immutable containers also reject type changes. + local_error = try + logdensity_and_gradient( + LogDensityFunction(local_model; adtype=AutoForwardDiff()), zeros(1) + ) + nothing + catch err + nested_exception(err) + end + @test local_error isa ArgumentError + + @model function indexed_value() + x = zeros(1) + return x[1] ~ Normal() + end + @model function property_value() + x = (a=0.0,) + return x.a ~ Normal() + end + value = ForwardDiff.Dual(0.1, 1.0) + for (model, vn) in + ((indexed_value(), @varname(x[1])), (property_value(), @varname(x.a))) + model = setthreadsafe(model, true) + for op in (condition, fix) + @test_throws ArgumentError init!!( + Xoshiro(1), + op(model, vn => value), + VarInfo(), + InitFromPrior(), + UnlinkAll(), + ) + end + end + end + @testset "assumes are threadsafe" begin # See https://github.com/TuringLang/DynamicPPL.jl/pull/1284. @model function threaded_assume() From bcb891ce15a738eda98b0c400e163f49dcd753e3 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 21 Oct 2026 14:51:13 +0100 Subject: [PATCH 2/7] docs: warn that thread-safe evaluation can reorder parameters Assisted-by: Codex --- docs/src/accs/threadsafe.md | 25 +++++++++++++++++-------- docs/src/ldf/models.md | 2 +- test/logdensityfunction.jl | 21 +++++++++++++++++++++ 3 files changed, 39 insertions(+), 9 deletions(-) diff --git a/docs/src/accs/threadsafe.md b/docs/src/accs/threadsafe.md index 41a20eb032..9107299e0a 100644 --- a/docs/src/accs/threadsafe.md +++ b/docs/src/accs/threadsafe.md @@ -76,6 +76,16 @@ Parameter values are also accumulator outputs. `VectorValueAccumulator` and This does not make arbitrary mutations in the model body or shared mutable initialisation strategies thread-safe. +Because task accumulators are combined in an unspecified order, the order of collected +variables can differ between evaluations, including variables defined outside the +threaded region. A `VarInfo` or `LogDensityFunction` built from such an evaluation +inherits that order. Each function is consistent with its own mapping, but using +`internal_values_as_vector` (`vi[:]`) from another `VarInfo` can silently give the wrong +density. Obtain vectors with that function's mapping using +[`VectorParamAccumulator`](@ref), or pass vectorised values with matching transforms to +[`to_vector_params`](@ref), as described in +[Obtaining new vectorised parameters](@ref ldf-vector-params). + ## Containers under thread-safe evaluation Under thread-safe evaluation, a tilde assignment throws an `ArgumentError` if it @@ -83,11 +93,10 @@ replaces a mutable container on its path or changes an immutable container's typ wherever the assignment occurs. Slice reads are checked by type because they return copies. Allocate with the parameter element type via `::Type{T}=Float64`: use `Vector{T}(undef, n)` for scalar tildes or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` -for vector-valued tildes. Assignments into data arguments that cannot be promoted also -throw, including `Vector{Missing}` and integer data imputed by a continuous -distribution; `Union{Missing,Float32}` data can still throw during initialisation from -the prior. Use `Union{Missing,Float64}` data for continuous imputation. The check does -not prevent lost writes from rebinding shared immutable values, including `NamedTuple`s -and immutable wrappers around mutable storage such as `OffsetArray`s of `Vector{Any}`, -or from concurrent writes through overlapping slices (`x[1:N][i] ~`). Serial evaluation -with thread-safe mode disabled retains its usual widening behaviour. +for vector-valued tildes. Under thread-safe evaluation with automatic differentiation, +imputing missing entries in a data argument now throws this error because the model +cannot reallocate the data argument to hold AD values. The check does not prevent lost +writes from rebinding shared immutable values, including `NamedTuple`s and immutable +wrappers around mutable storage such as `OffsetArray`s of `Vector{Any}`, or from +concurrent writes through overlapping slices (`x[1:N][i] ~`). Serial evaluation with +thread-safe mode disabled retains its usual widening behaviour. diff --git a/docs/src/ldf/models.md b/docs/src/ldf/models.md index a5108a2ce0..abff9914ca 100644 --- a/docs/src/ldf/models.md +++ b/docs/src/ldf/models.md @@ -87,7 +87,7 @@ This represents a loss of information since we no longer know the individual com It is much safer to use `InitFromVector`, which encapsulates all the information needed to rerun the model, but does not pretend to also contain other information that is not actually updated. -## Obtaining new vectorised parameters +## [Obtaining new vectorised parameters](@id ldf-vector-params) We can also do the reverse process here, which is to obtain new vectorised parameters that are consistent with this `LogDensityFunction`. diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 9d892a3313..7fb0c435e9 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -107,6 +107,27 @@ end end end +@testset "LogDensityFunction: encoding respects root order" begin + @model function f() + x ~ Normal() + return y ~ Normal(x) + end + model = f() + x = TransformedValue([0.1], Unlink()) + y = TransformedValue([0.5], Unlink()) + # Thread-safe evaluation can collect these variables in either root order. + values_xy = VarNamedTuple(; x, y) + values_yx = VarNamedTuple(; y, x) + ldf_xy = LogDensityFunction(model, getlogjoint_internal, values_xy) + ldf_yx = LogDensityFunction(model, getlogjoint_internal, values_yx) + vector_xy = to_vector_params(values_xy, ldf_xy) + vector_yx = to_vector_params(values_xy, ldf_yx) + + @test vector_xy != vector_yx + @test LogDensityProblems.logdensity(ldf_xy, vector_xy) ≈ + LogDensityProblems.logdensity(ldf_yx, vector_yx) +end + @testset "LogDensityFunction: accessor functions" begin @model f() = x ~ Normal() ldf = LogDensityFunction(f(), getlogprior, UnlinkAll()) From 4437623e391d6fb6755d1927b7e2ed323101e119 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Thu, 22 Oct 2026 10:31:12 +0530 Subject: [PATCH 3/7] Document the workarounds for the thread-safe container check --- HISTORY.md | 2 +- docs/src/accs/threadsafe.md | 10 ++++++---- src/threadsafe.jl | 5 +++-- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 4292074f58..2f12539329 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,6 @@ # 0.43.0 (unreleased) -Thread-safe evaluation now throws an `ArgumentError` when a tilde assignment replaces a mutable container or changes an immutable container's type; allocate mutable storage with `::Type{T}=Float64` and `Vector{T}(undef, n)`. +Thread-safe evaluation now throws an `ArgumentError` when a tilde assignment replaces a mutable container or changes an immutable container's type. Allocate mutable storage with `::Type{T}=Float64` and `Vector{T}(undef, n)`. Imputing missing entries in a data argument now throws this error under automatic differentiation, because the data argument cannot be reallocated to hold AD values. Pass that data as `Vector{Union{Missing,Real}}` to keep imputing. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). `VarNamedTuple` membership resolves `begin` and `end` against the stored array. See [#1490](https://github.com/TuringLang/DynamicPPL.jl/pull/1490). diff --git a/docs/src/accs/threadsafe.md b/docs/src/accs/threadsafe.md index 9107299e0a..cd8fffb9c4 100644 --- a/docs/src/accs/threadsafe.md +++ b/docs/src/accs/threadsafe.md @@ -92,10 +92,12 @@ Under thread-safe evaluation, a tilde assignment throws an `ArgumentError` if it replaces a mutable container on its path or changes an immutable container's type, wherever the assignment occurs. Slice reads are checked by type because they return copies. Allocate with the parameter element type via `::Type{T}=Float64`: use -`Vector{T}(undef, n)` for scalar tildes or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` -for vector-valued tildes. Under thread-safe evaluation with automatic differentiation, -imputing missing entries in a data argument now throws this error because the model -cannot reallocate the data argument to hold AD values. The check does not prevent lost +`Vector{T}(undef, n)` for scalar tildes, `Matrix{T}` with `x[:, i] ~ MvNormal(...)` +for vector-valued tildes, or `(a = zero(T),)` for a `NamedTuple`. Under thread-safe +evaluation with automatic differentiation, imputing missing entries in a data argument +now throws this error because the model cannot reallocate the data argument to hold AD +values. Pass that data as `Vector{Union{Missing,Real}}` to keep imputing, which gives +the same log density and gradient as serial evaluation. The check does not prevent lost writes from rebinding shared immutable values, including `NamedTuple`s and immutable wrappers around mutable storage such as `OffsetArray`s of `Vector{Any}`, or from concurrent writes through overlapping slices (`x[1:N][i] ~`). Serial evaluation with diff --git a/src/threadsafe.jl b/src/threadsafe.jl index d34751ffe4..78454b154e 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -59,10 +59,11 @@ end "Thread-safe evaluation cannot replace the container of tilde variable " * "$(vn): " * replacement * - "Mutable containers must retain their identity; immutable containers must " * + "Mutable containers must retain their identity. Immutable containers must " * "retain their type. Allocate with the parameter element type (for example, " * "`::Type{T}=Float64` with `Vector{T}(undef, n)` for scalar tildes, " * - "or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes).", + "`Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes, " * + "or `(a = zero(T),)` for a `NamedTuple`).", ), ) end From 6751dfd25cd19a5707f423d00a1f7732b3e665ea Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 22 Oct 2026 11:03:17 +0100 Subject: [PATCH 4/7] compiler: define the serial tilde setter with its call sites Assisted-by: Codex --- src/compiler.jl | 4 ++++ src/threadsafe.jl | 4 ---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index f3547e350e..5677e84bbc 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -426,6 +426,10 @@ function generate_tilde_literal(left, right) end end +function set_tilde_value!!(::AbstractVarInfo, object, value, vn::VarName) + return Accessors.set(object, AbstractPPL.with_mutation(AbstractPPL.getoptic(vn)), value) +end + assign_or_set!!(lhs::Symbol, rhs, vn) = drop_escape(:($lhs = $rhs)) function assign_or_set!!(lhs::Expr, rhs, vn) left_top_sym = get_top_level_symbol(lhs) diff --git a/src/threadsafe.jl b/src/threadsafe.jl index 78454b154e..f0a390350a 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -44,10 +44,6 @@ function _threadsafe_varinfo(vi::AbstractVarInfo, ::Type{L}) where {L<:Accumulat end ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi -function set_tilde_value!!(::AbstractVarInfo, object, value, vn::VarName) - return Accessors.set(object, AbstractPPL.with_mutation(AbstractPPL.getoptic(vn)), value) -end - @noinline function _tilde_container_error(old, new, vn) replacement = if typeof(old) === typeof(new) "$(typeof(old)) was replaced by a copy of the same type. " From d43f035abb65782f940ea78c03302cd750b3e59e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 22 Oct 2026 11:33:57 +0100 Subject: [PATCH 5/7] threadsafe: require tilde writes to keep containers Rebinding shared immutable values such as `x.a ~ Normal()` can lose writes even when their types stay unchanged. Require each container on a tilde path to remain identical. Traverse intermediate array slices through views and update only the leaf. Immutable replacements now throw even in task-local code without AD. Use mutable storage allocated with the parameter element type. Assisted-by: Codex --- HISTORY.md | 2 +- docs/src/accs/threadsafe.md | 35 +++++---- src/threadsafe.jl | 71 +++++++++--------- test/threadsafe.jl | 141 +++++++++++++++++++++++++++++++++--- 4 files changed, 189 insertions(+), 60 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 2f12539329..2cc48fd8c1 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,6 @@ # 0.43.0 (unreleased) -Thread-safe evaluation now throws an `ArgumentError` when a tilde assignment replaces a mutable container or changes an immutable container's type. Allocate mutable storage with `::Type{T}=Float64` and `Vector{T}(undef, n)`. Imputing missing entries in a data argument now throws this error under automatic differentiation, because the data argument cannot be reallocated to hold AD values. Pass that data as `Vector{Union{Missing,Real}}` to keep imputing. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). +Thread-safe tilde assignments now throw an `ArgumentError` if any container is replaced, including same-type immutable values: `(a = zero(T),)` → `Vector{T}(undef, n)` or `Ref(zero(T))`. Allocate mutable storage with `::Type{T}=Float64`; use `Vector{Union{Missing,Real}}` for imputation under AD. Intermediate slices now write through views. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). `VarNamedTuple` membership resolves `begin` and `end` against the stored array. See [#1490](https://github.com/TuringLang/DynamicPPL.jl/pull/1490). diff --git a/docs/src/accs/threadsafe.md b/docs/src/accs/threadsafe.md index cd8fffb9c4..c7a0214964 100644 --- a/docs/src/accs/threadsafe.md +++ b/docs/src/accs/threadsafe.md @@ -88,17 +88,24 @@ density. Obtain vectors with that function's mapping using ## Containers under thread-safe evaluation -Under thread-safe evaluation, a tilde assignment throws an `ArgumentError` if it -replaces a mutable container on its path or changes an immutable container's type, -wherever the assignment occurs. Slice reads are checked by type because they return -copies. Allocate with the parameter element type via `::Type{T}=Float64`: use -`Vector{T}(undef, n)` for scalar tildes, `Matrix{T}` with `x[:, i] ~ MvNormal(...)` -for vector-valued tildes, or `(a = zero(T),)` for a `NamedTuple`. Under thread-safe -evaluation with automatic differentiation, imputing missing entries in a data argument -now throws this error because the model cannot reallocate the data argument to hold AD -values. Pass that data as `Vector{Union{Missing,Real}}` to keep imputing, which gives -the same log density and gradient as serial evaluation. The check does not prevent lost -writes from rebinding shared immutable values, including `NamedTuple`s and immutable -wrappers around mutable storage such as `OffsetArray`s of `Vector{Any}`, or from -concurrent writes through overlapping slices (`x[1:N][i] ~`). Serial evaluation with -thread-safe mode disabled retains its usual widening behaviour. +Under thread-safe evaluation, every container on a tilde assignment's path must +retain its identity (`===`); otherwise the assignment throws an `ArgumentError`. +The written leaf is exempt. Intermediate array slices use views, so +`x[1:2][i] ~ Normal()` writes only element `i` and `x[:, j] .~ Normal()` writes +in place. A final slice, as in `x[:, j] ~ MvNormal(...)`, is written into its parent. +Concurrent writes to the same element still require synchronisation by the caller. + +Allocate mutable storage with the parameter element type via `::Type{T}=Float64`: +use `Vector{T}(undef, n)` for scalar tildes or `Matrix{T}` with +`x[:, i] ~ MvNormal(...)` for vector-valued tildes. Replace a `NamedTuple`, tuple, +or static array containing scalar parameters with a `Vector` or `Ref`, or use +plain scalar variables. This rule also applies to single-task and task-local +immutable containers. A `NamedTuple` of vectors can stay unchanged while its +vectors are updated in place. Immutable wrappers such as an `OffsetArray` must +also retain their identity; replacing their backing storage is rejected. + +Imputing missing entries in a data argument under automatic differentiation +requires storage that can hold AD values. Pass the data as +`Vector{Union{Missing,Real}}` to retain the same log density and gradient as serial +evaluation. Serial evaluation with thread-safe mode disabled retains its usual +widening behaviour. diff --git a/src/threadsafe.jl b/src/threadsafe.jl index f0a390350a..2c43a93291 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -55,54 +55,55 @@ ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi "Thread-safe evaluation cannot replace the container of tilde variable " * "$(vn): " * replacement * - "Mutable containers must retain their identity. Immutable containers must " * - "retain their type. Allocate with the parameter element type (for example, " * + "Thread-safe evaluation requires tilde variables to be written in place. " * + "Every container on the path must retain its identity (`===`). " * + "Allocate mutable storage with the parameter element type (for example, " * "`::Type{T}=Float64` with `Vector{T}(undef, n)` for scalar tildes, " * - "`Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes, " * - "or `(a = zero(T),)` for a `NamedTuple`).", + "or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes). " * + "Use a `Vector` or `Ref` instead of a `NamedTuple`, tuple or static array.", ), ) end -@inline function _capture_tilde_path(old, optic) - tail = AbstractPPL.otail(optic) - tail isa AbstractPPL.Iden && return ((old, false),) - sliced = _tilde_reads_slice(optic, old) - # The final slice copy is written back into its checked parent; reading it here - # would only allocate another copy. - if AbstractPPL.otail(tail) isa AbstractPPL.Iden && sliced - return ((old, sliced),) - end - child = AbstractPPL.ohead(optic)(old) - return ((old, sliced), _capture_tilde_path(child, tail)...) +@inline _tilde_child(object, optic) = optic(object) +@inline function _tilde_child(object::AbstractArray, optic::AbstractPPL.Index) + indices = AbstractPPL.concretize(optic, object) + return Base.maybeview(object, indices.ix...; indices.kw...) end -function _tilde_reads_slice(optic::AbstractPPL.Index, object::AbstractArray) - indices = AbstractPPL.concretize(AbstractPPL.ohead(optic), object) - return any( - i -> i isa Union{AbstractArray,Colon}, (indices.ix..., values(indices.kw)...) - ) +@inline function _set_tilde_leaf(object, value, optic, vn) + return Accessors.set(object, AbstractPPL.with_mutation(optic), value) end -_tilde_reads_slice(optic, object) = false -@inline function _check_tilde_path(old_path, new, optic, vn, sliced=false) - old, child_sliced = first(old_path) - valid = ismutable(old) && !sliced ? old === new : typeof(old) === typeof(new) - valid || _tilde_container_error(old, new, vn) +# BangBang has no zero-index setter for RefValue. +@inline function _set_tilde_leaf( + object::Base.RefValue, value, optic::AbstractPPL.Index{Tuple{}}, vn +) + try + return setindex!(object, value) + catch err + err isa Union{MethodError,InexactError} || rethrow() + return _tilde_container_error(object, Ref(value), vn) + end +end + +@inline function _set_tilde_path(object, value, optic, vn) tail = AbstractPPL.otail(optic) - tail isa AbstractPPL.Iden && return nothing - AbstractPPL.otail(tail) isa AbstractPPL.Iden && child_sliced && return nothing - return _check_tilde_path( - Base.tail(old_path), AbstractPPL.ohead(optic)(new), tail, vn, child_sliced - ) + if tail isa AbstractPPL.Iden + new = _set_tilde_leaf(object, value, optic, vn) + object === new || _tilde_container_error(object, new, vn) + else + head = AbstractPPL.ohead(optic) + child = _tilde_child(object, head) + _set_tilde_path(child, value, tail, vn) + new_child = _tilde_child(object, head) + child === new_child || _tilde_container_error(child, new_child, vn) + end + return object end function set_tilde_value!!(::ThreadSafeVarInfo, object, value, vn::VarName) - optic = AbstractPPL.getoptic(vn) - old_path = _capture_tilde_path(object, optic) - new = Accessors.set(object, AbstractPPL.with_mutation(optic), value) - _check_tilde_path(old_path, new, optic, vn) - return new + return _set_tilde_path(object, value, AbstractPPL.getoptic(vn), vn) end function _get_task_accs_cache(vi::ThreadSafeVarInfo{V,L}) where {V,L} diff --git a/test/threadsafe.jl b/test/threadsafe.jl index ac39fa4b08..72b4b76604 100644 --- a/test/threadsafe.jl +++ b/test/threadsafe.jl @@ -8,6 +8,7 @@ using Distributions using DynamicPPL using Test using ForwardDiff +using OffsetArrays: OffsetArray using Random: Xoshiro using ADTypes: AutoForwardDiff using LogDensityProblems: logdensity, logdensity_and_gradient @@ -567,18 +568,138 @@ end ) ≈ [sum(y)] local_model = setthreadsafe(local_values(), true) - local_ldf = LogDensityFunction(local_model) - @test logdensity(local_ldf, zeros(1)) ≈ 8 * logpdf(Normal(), 0.0) - # Accepted limitation: task-local immutable containers also reject type changes. - local_error = try - logdensity_and_gradient( - LogDensityFunction(local_model; adtype=AutoForwardDiff()), zeros(1) + for adtype in (nothing, AutoForwardDiff()) + local_error = try + LogDensityFunction(local_model; adtype=adtype) + nothing + catch err + nested_exception(err) + end + @test local_error isa ArgumentError + end + + @model function immutable_value() + x = (a=0.0,) + x.a ~ Normal() + return x.a + end + @model function nested_immutable_value() + x = Box((a=0.0,)) + x.v.a ~ Normal() + return x.v.a + end + @model function intermediate_slice() + x = zeros(2) + x[1:2][1] ~ Normal() + return x[1] + end + @model function tuple_value() + x = (0.0,) + x[1] ~ Normal() + return x[1] + end + @model function offset_value() + x = OffsetArray(Any[zeros(1)], 0:0) + x[0] ~ MvNormal(zeros(1), ones(1)) + return only(x[0]) + end + for (model, params) in ( + (immutable_value(), (; x=(a=0.25,))), + (nested_immutable_value(), (; x=Box((a=0.25,)))), + (tuple_value(), (; x=[0.25])), + (offset_value(), (; x=OffsetArray([[0.25]], 0:0))), + ) + strategy = InitFromParams(params, nothing) + @test first(init!!(model, VarInfo(), strategy, UnlinkAll())) == 0.25 + @test_throws ArgumentError init!!( + setthreadsafe(model, true), VarInfo(), strategy, UnlinkAll() + ) + end + for op in (condition, fix) + @test_throws ArgumentError init!!( + op(setthreadsafe(immutable_value(), true), @varname(x.a) => 0.25), + VarInfo(), + InitFromPrior(), + UnlinkAll(), + ) + end + + @model function immutable_vectors(::Type{T}=Float64) where {T} + x = (a=zeros(T, 1),) + x.a[1] ~ Normal() + return x.a[1] + end + @model function reference_value(::Type{T}=Float64) where {T} + r = Ref(zero(T)) + r[] ~ Normal() + return r[] + end + @model function mutable_field(::Type{T}=Float64) where {T} + s = Box(zero(T)) + s.v ~ Normal() + return s.v + end + @model function dictionary_value(::Type{T}=Float64) where {T} + key = [1] + x = Dict(key => zeros(T, 1)) + x[key][1] ~ Normal() + return x[key][1] + end + for model in + (immutable_vectors(), reference_value(), mutable_field(), dictionary_value()) + result, vi = init!!( + Xoshiro(1), + setthreadsafe(model, true), + VarInfo(), + InitFromPrior(), + UnlinkAll(), ) - nothing - catch err - nested_exception(err) + expected = rand(Xoshiro(1), Normal()) + @test result == expected + @test getlogjoint(vi) ≈ logpdf(Normal(), expected) + end + for (model, params) in ( + (immutable_vectors(), z -> (; x=(a=[z],))), + (reference_value(), z -> (; r=fill(z))), + (mutable_field(), z -> (; s=Box(z))), + ) + model = setthreadsafe(model, true) + @test ForwardDiff.derivative(z -> logjoint(model, params(z)), 0.25) ≈ -0.25 + @test ForwardDiff.derivative(0.25) do z + first( + init!!( + model, VarInfo(), InitFromParams(params(z), nothing), UnlinkAll() + ), + ) + end ≈ 1.0 end - @test local_error isa ArgumentError + @test first( + init!!( + setthreadsafe(intermediate_slice(), true), + VarInfo(), + InitFromParams((; x=[0.25, 0.25]), nothing), + UnlinkAll(), + ), + ) == 0.25 + @test_throws ArgumentError init!!( + setthreadsafe(intermediate_slice(), true), + VarInfo(), + InitFromParams((; x=[ForwardDiff.Dual(0.25, 1.0), 0.25]), nothing), + UnlinkAll(), + ) + + @model function impute(y) + for i in eachindex(y) + y[i] ~ Normal() + end + end + imputed_density(z, threadsafe) = logjoint( + setthreadsafe(impute(Union{Missing,Real}[0.5, missing]), threadsafe), + (; y=[0.5, z]), + ) + @test imputed_density(0.25, true) ≈ imputed_density(0.25, false) + @test imputed_density(0.25, true) ≈ logpdf(Normal(), 0.5) + logpdf(Normal(), 0.25) + @test ForwardDiff.derivative(z -> imputed_density(z, true), 0.25) ≈ -0.25 @model function indexed_value() x = zeros(1) From 2b2cc1bd20ab32cd83cab8eed52336cc32ab9e79 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 22 Oct 2026 12:17:29 +0100 Subject: [PATCH 6/7] history: drop mechanism from thread-safe entry Assisted-by: Codex --- HISTORY.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index 2cc48fd8c1..f51e8eaac6 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,6 @@ # 0.43.0 (unreleased) -Thread-safe tilde assignments now throw an `ArgumentError` if any container is replaced, including same-type immutable values: `(a = zero(T),)` → `Vector{T}(undef, n)` or `Ref(zero(T))`. Allocate mutable storage with `::Type{T}=Float64`; use `Vector{Union{Missing,Real}}` for imputation under AD. Intermediate slices now write through views. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). +Thread-safe tilde assignments now throw an `ArgumentError` if any container is replaced, including same-type immutable values: `(a = zero(T),)` → `Vector{T}(undef, n)` or `Ref(zero(T))`. Allocate mutable storage with `::Type{T}=Float64`; use `Vector{Union{Missing,Real}}` for imputation under AD. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). `VarNamedTuple` membership resolves `begin` and `end` against the stored array. See [#1490](https://github.com/TuringLang/DynamicPPL.jl/pull/1490). From a592b3b00024de6a1601dfbadc18a584d6a823d0 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 22 Oct 2026 21:51:57 +0100 Subject: [PATCH 7/7] threadsafe: stop recommending Ref storage Assisted-by: Codex --- HISTORY.md | 2 +- docs/src/accs/threadsafe.md | 2 +- src/threadsafe.jl | 14 +------------- test/threadsafe.jl | 12 ++---------- 4 files changed, 5 insertions(+), 25 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index f51e8eaac6..f7be19faad 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,6 @@ # 0.43.0 (unreleased) -Thread-safe tilde assignments now throw an `ArgumentError` if any container is replaced, including same-type immutable values: `(a = zero(T),)` → `Vector{T}(undef, n)` or `Ref(zero(T))`. Allocate mutable storage with `::Type{T}=Float64`; use `Vector{Union{Missing,Real}}` for imputation under AD. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). +Thread-safe tilde assignments now throw an `ArgumentError` if any container is replaced, including same-type immutable values: `(a = zero(T),)` → `Vector{T}(undef, n)`. Allocate mutable storage with `::Type{T}=Float64`; use `Vector{Union{Missing,Real}}` for imputation under AD. See [#1521](https://github.com/TuringLang/DynamicPPL.jl/pull/1521). `VarNamedTuple` membership resolves `begin` and `end` against the stored array. See [#1490](https://github.com/TuringLang/DynamicPPL.jl/pull/1490). diff --git a/docs/src/accs/threadsafe.md b/docs/src/accs/threadsafe.md index c7a0214964..f62f82e99e 100644 --- a/docs/src/accs/threadsafe.md +++ b/docs/src/accs/threadsafe.md @@ -98,7 +98,7 @@ Concurrent writes to the same element still require synchronisation by the calle Allocate mutable storage with the parameter element type via `::Type{T}=Float64`: use `Vector{T}(undef, n)` for scalar tildes or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes. Replace a `NamedTuple`, tuple, -or static array containing scalar parameters with a `Vector` or `Ref`, or use +or static array containing scalar parameters with a `Vector`, or use plain scalar variables. This rule also applies to single-task and task-local immutable containers. A `NamedTuple` of vectors can stay unchanged while its vectors are updated in place. Immutable wrappers such as an `OffsetArray` must diff --git a/src/threadsafe.jl b/src/threadsafe.jl index 2c43a93291..cca58a9f23 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -60,7 +60,7 @@ ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi "Allocate mutable storage with the parameter element type (for example, " * "`::Type{T}=Float64` with `Vector{T}(undef, n)` for scalar tildes, " * "or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes). " * - "Use a `Vector` or `Ref` instead of a `NamedTuple`, tuple or static array.", + "Use a `Vector` instead of a `NamedTuple`, tuple or static array.", ), ) end @@ -75,18 +75,6 @@ end return Accessors.set(object, AbstractPPL.with_mutation(optic), value) end -# BangBang has no zero-index setter for RefValue. -@inline function _set_tilde_leaf( - object::Base.RefValue, value, optic::AbstractPPL.Index{Tuple{}}, vn -) - try - return setindex!(object, value) - catch err - err isa Union{MethodError,InexactError} || rethrow() - return _tilde_container_error(object, Ref(value), vn) - end -end - @inline function _set_tilde_path(object, value, optic, vn) tail = AbstractPPL.otail(optic) if tail isa AbstractPPL.Iden diff --git a/test/threadsafe.jl b/test/threadsafe.jl index 72b4b76604..5a09656a06 100644 --- a/test/threadsafe.jl +++ b/test/threadsafe.jl @@ -629,11 +629,6 @@ end x.a[1] ~ Normal() return x.a[1] end - @model function reference_value(::Type{T}=Float64) where {T} - r = Ref(zero(T)) - r[] ~ Normal() - return r[] - end @model function mutable_field(::Type{T}=Float64) where {T} s = Box(zero(T)) s.v ~ Normal() @@ -645,8 +640,7 @@ end x[key][1] ~ Normal() return x[key][1] end - for model in - (immutable_vectors(), reference_value(), mutable_field(), dictionary_value()) + for model in (immutable_vectors(), mutable_field(), dictionary_value()) result, vi = init!!( Xoshiro(1), setthreadsafe(model, true), @@ -659,9 +653,7 @@ end @test getlogjoint(vi) ≈ logpdf(Normal(), expected) end for (model, params) in ( - (immutable_vectors(), z -> (; x=(a=[z],))), - (reference_value(), z -> (; r=fill(z))), - (mutable_field(), z -> (; s=Box(z))), + (immutable_vectors(), z -> (; x=(a=[z],))), (mutable_field(), z -> (; s=Box(z))) ) model = setthreadsafe(model, true) @test ForwardDiff.derivative(z -> logjoint(model, params(z)), 0.25) ≈ -0.25