fix merge

Author: penelopeysm

src/model.jl

@@ -937,10 +937,7 @@ Returns a tuple of the model's return value, plus the updated `varinfo` object.
 )
     ctx = InitContext(rng, strategy)
     model = DynamicPPL.setleafcontext(model, ctx)
-    # TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
-    # it _should_ do, but this is wrong regardless.
-    # https://github.com/TuringLang/DynamicPPL.jl/issues/1086
-    return if Threads.nthreads() > 1
+    return if _requires_threadsafe(model)
         # TODO(penelopeysm): The logic for setting eltype of accs is very similar to that
         # used in `unflatten`. The reason why we need it here is because the VarInfo `vi`
         # won't have been filled with parameters prior to `init!!` being called.

test/logdensityfunction.jl

@@ -51,21 +51,19 @@ using Mooncake: Mooncake
     end
 
     @testset "Threaded observe" begin
-        if Threads.nthreads() > 1
-            @model function threaded(y)
-                x ~ Normal()
-                Threads.@threads for i in eachindex(y)
-                    y[i] ~ Normal(x)
-                end
+        @model function threaded(y)
+            x ~ Normal()
+            Threads.@threads for i in eachindex(y)
+                y[i] ~ Normal(x)
             end
-            N = 100
-            model = threaded(zeros(N))
-            ldf = DynamicPPL.LogDensityFunction(model)
-
-            xs = [1.0]
-            @test LogDensityProblems.logdensity(ldf, xs) ≈
-                logpdf(Normal(), xs[1]) + N * logpdf(Normal(xs[1]), 0.0)
         end
+        N = 100
+        model = setthreadsafe(threaded(zeros(N)), true)
+        ldf = DynamicPPL.LogDensityFunction(model)
+
+        xs = [1.0]
+        @test LogDensityProblems.logdensity(ldf, xs) ≈
+            logpdf(Normal(), xs[1]) + N * logpdf(Normal(xs[1]), 0.0)
     end
 end
 
@@ -109,34 +107,32 @@ end
 end
 
 @testset "LogDensityFunction: performance" begin
-    if Threads.nthreads() == 1
-        # Evaluating these three models should not lead to any allocations (but only when
-        # not using TSVI).
-        @model function f()
-            x ~ Normal()
-            return 1.0 ~ Normal(x)
-        end
-        @model function submodel_inner()
-            m ~ Normal(0, 1)
-            s ~ Exponential()
-            return (m=m, s=s)
-        end
-        # Note that for the allocation tests to work on this one, `inner` has
-        # to be passed as an argument to `submodel_outer`, instead of just
-        # being called inside the model function itself
-        @model function submodel_outer(inner)
-            params ~ to_submodel(inner)
-            y ~ Normal(params.m, params.s)
-            return 1.0 ~ Normal(y)
-        end
-        @testset for model in
-                     (f(), submodel_inner() | (; s=0.0), submodel_outer(submodel_inner()))
-            vi = VarInfo(model)
-            ldf = DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi)
-            x = vi[:]
-            bench = median(@be LogDensityProblems.logdensity(ldf, x))
-            @test iszero(bench.allocs)
-        end
+    # Evaluating these three models should not lead to any allocations (but only when
+    # not using TSVI).
+    @model function f()
+        x ~ Normal()
+        return 1.0 ~ Normal(x)
+    end
+    @model function submodel_inner()
+        m ~ Normal(0, 1)
+        s ~ Exponential()
+        return (m=m, s=s)
+    end
+    # Note that for the allocation tests to work on this one, `inner` has
+    # to be passed as an argument to `submodel_outer`, instead of just
+    # being called inside the model function itself
+    @model function submodel_outer(inner)
+        params ~ to_submodel(inner)
+        y ~ Normal(params.m, params.s)
+        return 1.0 ~ Normal(y)
+    end
+    @testset for model in
+                 (f(), submodel_inner() | (; s=0.0), submodel_outer(submodel_inner()))
+        vi = VarInfo(model)
+        ldf = DynamicPPL.LogDensityFunction(model, DynamicPPL.getlogjoint_internal, vi)
+        x = vi[:]
+        bench = median(@be LogDensityProblems.logdensity(ldf, x))
+        @test iszero(bench.allocs)
     end
 end