update setparams!!

TuringLang · Oct 22, 2024 · 4b0be7f · 4b0be7f
1 parent 933efe5
commit 4b0be7f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/research/tests/runtests.jl b/research/tests/runtests.jl
@@ -11,6 +11,6 @@ include("../src/riemannian_hmc.jl")
 include("relativistic_hmc.jl")
 include("riemannian_hmc.jl")
 
-@main function runtests(patterns...; dry::Bool = false)
+Comonicon.@main function runtests(patterns...; dry::Bool = false)
     retest(patterns...; dry = dry, verbose = Inf)
 end
diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl
@@ -35,8 +35,12 @@ function AbstractMCMC.getparams(state::HMCState)
     return state.transition.z.θ
 end
 
-function AbstractMCMC.setparams!!(state::HMCState, θ)
-    return @set state.transition.z.θ = θ
+function AbstractMCMC.setparams!!(state::HMCState, params)
+    hamiltonian = AdvancedHMC.Hamiltonian(state.metric, model)
+    return Setfield.@set state.transition.z = AdvancedHMC.phasepoint(
+        hamiltonian, params, state.transition.z.r;
+        ℓκ=state.transition.z.ℓκ
+    )
 end
 
 """