Transition (#2026)

* 1st commit

* should be working

* connector to old struct

* different type

* bug

* bug

* Update src/inference/Inference.jl

David's suggestion

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Update src/inference/Inference.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* remove reference to gibbs transitions int tests

* Update src/inference/Inference.jl

---------

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Author: 3 people

src/inference/Inference.jl

@@ -105,18 +105,28 @@ end
 # Default Transition #
 ######################
 
-struct Transition{T, F<:AbstractFloat}
-    θ  :: T
-    lp :: F
+struct Transition{T, F<:AbstractFloat, S<:Union{NamedTuple, Nothing}}
+    θ     :: T
+    lp    :: F # TODO: merge `lp` with `stat`
+    stat  :: S
 end
 
-function Transition(vi::AbstractVarInfo, nt::NamedTuple=NamedTuple())
-    theta = merge(tonamedtuple(vi), nt)
+Transition(θ, lp) = Transition(θ, lp, nothing)
+
+function Transition(vi::AbstractVarInfo; nt::NamedTuple=NamedTuple())
+    θ = merge(tonamedtuple(vi), nt)
     lp = getlogp(vi)
-    return Transition{typeof(theta), typeof(lp)}(theta, lp)
+    return Transition(θ, lp, nothing)
 end
 
-metadata(t::Transition) = (lp = t.lp,)
+function metadata(t::Transition)
+    stat = t.stat
+    if stat === nothing
+        return (lp = t.lp,)
+    else
+        return merge((lp = t.lp,), stat)
+    end
+end
 
 DynamicPPL.getlogp(t::Transition) = t.lp
 

src/inference/gibbs.jl

@@ -69,23 +69,6 @@ struct GibbsState{V<:VarInfo,S<:Tuple{Vararg{Sampler}},T}
     states::T
 end
 
-struct GibbsTransition{T,F}
-    "The parameters for any given sample."
-    θ::T
-    "The joint log probability for the sample's parameters."
-    lp::F
-end
-
-function GibbsTransition(vi::AbstractVarInfo)
-    theta = tonamedtuple(vi)
-    lp = getlogp(vi)
-    return GibbsTransition(theta, lp)
-end
-
-metadata(t::GibbsTransition) = (lp = t.lp,)
-
-DynamicPPL.getlogp(t::GibbsTransition) = t.lp
-
 # extract varinfo object from state
 """
     gibbs_varinfo(model, sampler, state)
@@ -213,7 +196,7 @@ function DynamicPPL.initialstep(
     end
 
     # Compute initial transition and state.
-    transition = GibbsTransition(vi)
+    transition = Transition(vi)
     state = GibbsState(vi, samplers, states)
 
     return transition, state
@@ -248,5 +231,5 @@ function AbstractMCMC.step(
         return newstate
     end
 
-    return GibbsTransition(vi), GibbsState(vi, samplers, states)
+    return Transition(vi), GibbsState(vi, samplers, states)
 end

src/inference/hmc.jl

@@ -21,24 +21,12 @@ end
 # Hamiltonian Transition #
 ##########################
 
-struct HMCTransition{T,NT<:NamedTuple,F<:AbstractFloat}
-    θ::T
-    lp::F
-    stat::NT
-end
-
-function HMCTransition(vi::AbstractVarInfo, t::AHMC.Transition)
+function Transition(vi::AbstractVarInfo, t::AHMC.Transition)
     theta = tonamedtuple(vi)
     lp = getlogp(vi)
-    return HMCTransition(theta, lp, t.stat)
+    return Transition(theta, lp, t.stat)
 end
 
-function metadata(t::HMCTransition)
-    return merge((lp = t.lp,), t.stat)
-end
-
-DynamicPPL.getlogp(t::HMCTransition) = t.lp
-
 ###
 ### Hamiltonian Monte Carlo samplers.
 ###
@@ -230,7 +218,7 @@ function DynamicPPL.initialstep(
         vi = setlogp!!(vi, log_density_old)
     end
 
-    transition = HMCTransition(vi, t)
+    transition = Transition(vi, t)
     state = HMCState(vi, 1, kernel, hamiltonian, t.z, adaptor)
 
     return transition, state
@@ -270,7 +258,7 @@ function AbstractMCMC.step(
     end
 
     # Compute next transition and state.
-    transition = HMCTransition(vi, t)
+    transition = Transition(vi, t)
     newstate = HMCState(vi, i, kernel, hamiltonian, t.z, state.adaptor)
 
     return transition, newstate

test/inference/gibbs.jl

@@ -74,13 +74,13 @@
             ::Type{MCMCChains.Chains};
             kwargs...
         )
-            samples isa Vector{<:Inference.GibbsTransition} ||
+            samples isa Vector{<:Inference.Transition} ||
                 error("incorrect transitions")
             return
         end
 
         function callback(rng, model, sampler, sample, state, i; kwargs...)
-            sample isa Inference.GibbsTransition || error("incorrect sample")
+            sample isa Inference.Transition || error("incorrect sample")
             return
         end