stop using invokelatest

Author: sunxd3

JuliaBUGS/src/model/abstractppl.jl

@@ -698,8 +698,14 @@ function evaluate!!(
     temperature=1.0,
     transformed=model.transformed,
 )
-    evaluation_env, log_densities = evaluate_with_values!!(
-        model, flattened_values; temperature=temperature, transformed=transformed
-    )
+    if model.evaluation_mode isa UseAutoMarginalization
+        evaluation_env, log_densities = evaluate_with_marginalization_values!!(
+            model, flattened_values; temperature=temperature, transformed=transformed
+        )
+    else
+        evaluation_env, log_densities = evaluate_with_values!!(
+            model, flattened_values; temperature=temperature, transformed=transformed
+        )
+    end
     return evaluation_env, log_densities.tempered_logjoint
 end

JuliaBUGS/src/model/bugsmodel.jl

@@ -643,7 +643,7 @@ function getparams(model::BUGSModel, evaluation_env=model.evaluation_env)
             end
         else
             (; node_function, loop_vars) = model.g[v]
-            dist = Base.invokelatest(node_function, evaluation_env, loop_vars)
+            dist = node_function(evaluation_env, loop_vars)
             transformed_value = Bijectors.transform(
                 Bijectors.bijector(dist), AbstractPPL.get(evaluation_env, v)
             )
@@ -679,7 +679,7 @@ function getparams(
             d[v] = value
         else
             (; node_function, loop_vars) = model.g[v]
-            dist = Base.invokelatest(node_function, evaluation_env, loop_vars)
+            dist = node_function(evaluation_env, loop_vars)
             d[v] = Bijectors.transform(Bijectors.bijector(dist), value)
         end
     end

JuliaBUGS/src/model/evaluation.jl

@@ -500,7 +500,7 @@ function _marginalize_recursive(
 
     if !is_stochastic
         # Deterministic node
-        value = Base.invokelatest(node_function, env, loop_vars)
+        value = node_function(env, loop_vars)
         new_env = BangBang.setindex!!(env, value, current_vn)
         result = _marginalize_recursive(
             model,
@@ -515,7 +515,7 @@ function _marginalize_recursive(
 
     elseif is_observed
         # Observed stochastic node
-        dist = Base.invokelatest(node_function, env, loop_vars)
+        dist = node_function(env, loop_vars)
         obs_value = AbstractPPL.get(env, current_vn)
         obs_logp = logpdf(dist, obs_value)
 
@@ -538,7 +538,7 @@ function _marginalize_recursive(
 
     elseif is_discrete_finite
         # Discrete finite unobserved node - marginalize out
-        dist = Base.invokelatest(node_function, env, loop_vars)
+        dist = node_function(env, loop_vars)
         possible_values = enumerate_discrete_values(dist)
 
         logp_branches = Vector{typeof(zero(eltype(parameter_values)))}(
@@ -571,7 +571,7 @@ function _marginalize_recursive(
 
     else
         # Continuous or discrete infinite unobserved node - use parameter values
-        dist = Base.invokelatest(node_function, env, loop_vars)
+        dist = node_function(env, loop_vars)
         b = Bijectors.bijector(dist)
 
         if !haskey(var_lengths, current_vn)

JuliaBUGS/test/model/evaluation.jl

@@ -18,18 +18,16 @@ function test_bugs_model_log_density(
     transformed_model = compile(model_def, data, inits)
     untransformed_model = JuliaBUGS.settrans(transformed_model, false)
 
-    # Allow world age to advance by calling the functions in a separate evaluation
-    Base.invokelatest() do
-        @test _logjoint(untransformed_model) ≈ expected_untransformed rtol = 1E-6
-        @test _logjoint(transformed_model) ≈ expected_transformed rtol = 1E-6
-
-        @test LogDensityProblems.logdensity(
-            transformed_model, JuliaBUGS.getparams(transformed_model)
-        ) ≈ expected_transformed rtol = 1E-6
-        @test LogDensityProblems.logdensity(
-            untransformed_model, JuliaBUGS.getparams(untransformed_model)
-        ) ≈ expected_untransformed rtol = 1E-6
-    end
+    # Evaluate directly; model compilation happens before these calls
+    @test _logjoint(untransformed_model) ≈ expected_untransformed rtol = 1E-6
+    @test _logjoint(transformed_model) ≈ expected_transformed rtol = 1E-6
+
+    @test LogDensityProblems.logdensity(
+        transformed_model, JuliaBUGS.getparams(transformed_model)
+    ) ≈ expected_transformed rtol = 1E-6
+    @test LogDensityProblems.logdensity(
+        untransformed_model, JuliaBUGS.getparams(untransformed_model)
+    ) ≈ expected_untransformed rtol = 1E-6
 end
 
 @testset "evaluate_with_rng!! - controlling sampling behavior" begin