Implement Turing.Inference.getlogp_external (#39)

* Implement Turing.Inference.getlogp_external

* Add tests

* Fix tests

* Don't set arch=x64 for macos-latest

Author: penelopeysm

.github/workflows/CI.yml

@@ -25,14 +25,11 @@ jobs:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
-        arch:
-          - x64
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1

Project.toml

@@ -1,6 +1,6 @@
 name = "SliceSampling"
 uuid = "43f4d3e8-9711-4a8c-bd1b-03ac73a255cf"
-version = "0.7.6"
+version = "0.7.7"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -21,7 +21,7 @@ Distributions = "0.25"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 Random = "1"
-Turing = "0.37, 0.38, 0.39"
+Turing = "0.39.5"
 julia = "1.10"
 
 [extras]

ext/SliceSamplingTuringExt.jl

@@ -22,30 +22,29 @@ Turing.Inference.isgibbscomponent(::SliceSampling.Slice) = true
 Turing.Inference.isgibbscomponent(::SliceSampling.SliceSteppingOut) = true
 Turing.Inference.isgibbscomponent(::SliceSampling.SliceDoublingOut) = true
 
-function Turing.Inference.getparams(
-    ::Turing.DynamicPPL.Model, sample::SliceSampling.UnivariateSliceState
-)
+const SliceSamplingStates = Union{
+    SliceSampling.UnivariateSliceState,
+    SliceSampling.GibbsState,
+    SliceSampling.HitAndRunState,
+    SliceSampling.LatentSliceState,
+    SliceSampling.GibbsPolarSliceState,
+}
+function Turing.Inference.getparams(::Turing.DynamicPPL.Model, sample::SliceSamplingStates)
     return sample.transition.params
 end
 
-function Turing.Inference.getparams(
-    ::Turing.DynamicPPL.Model, state::SliceSampling.GibbsState
-)
-    return state.transition.params
-end
-
-function Turing.Inference.getparams(
-    ::Turing.DynamicPPL.Model, state::SliceSampling.HitAndRunState
+function Turing.Inference.getlogp_external(
+    ::Turing.DynamicPPL.Model, t::SliceSampling.Transition, state
 )
-    return state.transition.params
+    return t.lp
 end
 # end
 
 function SliceSampling.initial_sample(rng::Random.AbstractRNG, ℓ::Turing.LogDensityFunction)
     model  = ℓ.model
     vi     = Turing.DynamicPPL.VarInfo(rng, model, Turing.SampleFromUniform())
     vi_spl = last(Turing.DynamicPPL.evaluate!!(model, rng, vi, Turing.SampleFromUniform()))
-    θ     = vi_spl[:]
+    θ      = vi_spl[:]
 
     init_attempt_count = 1
     while !all(isfinite.(θ))

test/turing.jl

@@ -8,8 +8,13 @@
         return nothing
     end
 
+    @model function logp_check()
+        a ~ Normal()
+        return b ~ Normal()
+    end
+
     n_samples = 1000
-    model     = demo()
+    model = demo()
 
     @testset for sampler in [
         RandPermGibbs(Slice(1)),
@@ -30,6 +35,15 @@
         )
 
         chain = sample(model, externalsampler(sampler), n_samples; progress=false)
+
+        chain_logp_check = sample(
+            logp_check(), externalsampler(sampler), 100; progress=false
+        )
+        @test isapprox(
+            logpdf.(Normal(), chain_logp_check[:a]) .+
+            logpdf.(Normal(), chain_logp_check[:b]),
+            chain_logp_check[:lp],
+        )
     end
 
     @testset "gibbs($sampler)" for sampler in [
@@ -46,5 +60,17 @@
             n_samples;
             progress=false,
         )
+
+        chain_logp_check = sample(
+            logp_check(),
+            Turing.Gibbs(:a => externalsampler(sampler), :b => externalsampler(sampler)),
+            100;
+            progress=false,
+        )
+        @test isapprox(
+            logpdf.(Normal(), chain_logp_check[:a]) .+
+            logpdf.(Normal(), chain_logp_check[:b]),
+            chain_logp_check[:lp],
+        )
     end
 end