Reject tilde assignments that replace containers under thread-safe evaluation

HISTORY.md

@@ -1,5 +1,7 @@
 # 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)`. 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).
 
 `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).

docs/src/accs/threadsafe.md

@@ -75,3 +75,37 @@ 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.
+
+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, 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 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.

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`.
 

src/compiler.jl

@@ -426,15 +426,17 @@ 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)
     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 +468,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 +563,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

src/threadsafe.jl

@@ -44,6 +44,56 @@ function _threadsafe_varinfo(vi::AbstractVarInfo, ::Type{L}) where {L<:Accumulat
 end
 ThreadSafeVarInfo(vi::ThreadSafeVarInfo) = vi
 
+@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 *
+            "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, " *
+            "or `Matrix{T}` with `x[:, i] ~ MvNormal(...)` for vector-valued tildes). " *
+            "Use a `Vector` instead of a `NamedTuple`, tuple or static array.",
+        ),
+    )
+end
+
+@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
+
+@inline function _set_tilde_leaf(object, value, optic, vn)
+    return Accessors.set(object, AbstractPPL.with_mutation(optic), value)
+end
+
+@inline function _set_tilde_path(object, value, optic, vn)
+    tail = AbstractPPL.otail(optic)
+    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)
+    return _set_tilde_path(object, value, AbstractPPL.getoptic(vn), vn)
+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.

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