fix test errors

TuringLang · Oct 31, 2024 · e902811 · e902811
1 parent 58162a3
commit e902811
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 10 deletions.
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
@@ -146,9 +146,13 @@ function AbstractMCMC.getparams(t::Transition)
     return t.params
 end
 
-function AbstractMCMC.setparams!!(model::AbstractMCMC.LogDensityModel, t::Transition, params)
-    t = BangBang.setproperty!!(t, :params, params)
-    return BangBang.setproperty!!(t, :lp, LogDensityProblems.logdensity(model.logdensity, params))
+# TODO (sunxd): remove `DensityModel` in favor of `AbstractMCMC.LogDensityModel`
+function AbstractMCMC.setparams!!(model::DensityModelOrLogDensityModel, t::Transition, params)
+    return Transition(
+        params,
+        logdensity(model, params),
+        t.accepted
+    )
 end
 
 # Include inference methods.

diff --git a/src/MALA.jl b/src/MALA.jl
@@ -24,10 +24,12 @@ function AbstractMCMC.getparams(t::GradientTransition)
     return t.params
 end
 
-function AbstractMCMC.setparams!!(model::AbstractMCMC.LogDensityModel, t::GradientTransition, params)
-    return AdvancedMH.GradientTransition(
+function AbstractMCMC.setparams!!(model::DensityModelOrLogDensityModel, t::GradientTransition, params)
+    lp, gradient = logdensity_and_gradient(model, params)
+    return GradientTransition(
         params,
-        AdvancedMH.logdensity_and_gradient(model.logdensity, params)...,
+        lp,
+        gradient,
         t.accepted
     )
 end

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -38,10 +38,18 @@ include("util.jl")
         t1, _ = AbstractMCMC.step(Random.default_rng(), model, StaticMH([Normal(0, 1), Normal(0, 1)]))
         t2, _ = AbstractMCMC.step(Random.default_rng(), model, MALA(x -> MvNormal(x, I)); initial_params=ones(2))
         for t in [t1, t2]
-            @test AbstractMCMC.getparams(t) == t.params
-            @test AbstractMCMC.setparams!!(model, t, AbstractMCMC.getparams(t)) == t
-            t_replaced = AbstractMCMC.setparams!!(model, t, (μ=1.0, σ=2.0))
-            @test t_replaced.params == (μ=1.0, σ=2.0)
+            @test AbstractMCMC.getparams(model, t) == t.params
+
+            new_transition = AbstractMCMC.setparams!!(model, t, AbstractMCMC.getparams(model, t))
+            @test new_transition.lp == t.lp
+            @test new_transition.accepted == t.accepted
+            @test new_transition.params == t.params
+            if hasfield(typeof(t), :gradient)
+                @test new_transition.gradient == t.gradient
+            end
+
+            t_replaced = AbstractMCMC.setparams!!(model, t, [1.0, 2.0])
+            @test t_replaced.params == [1.0, 2.0]
         end
     end