Proper type dispatch (#91)

* Proper type dispatch * Bump
TuringLang · Nov 4, 2023 · 8fbcb20 · 8fbcb20 · yebai · Nov 4, 2023
1 parent 3700563
commit 8fbcb20
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "AdvancedPS"
 uuid = "576499cb-2369-40b2-a588-c64705576edc"
 authors = ["TuringLang"]
-version = "0.5.2"
+version = "0.5.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

diff --git a/src/container.jl b/src/container.jl
@@ -62,8 +62,8 @@ end
 Update reference trajectory. Defaults to `nothing`
 """
 function update_ref!(
-    particle::Trace, pc::ParticleContainer, sampler::AbstractParticleSampler
-)
+    particle::Trace, pc::ParticleContainer, sampler::T
+) where {T<:AbstractMCMC.AbstractSampler}
     return nothing
 end
 
@@ -171,11 +171,11 @@ of the particle `weights`. For Particle Gibbs sampling, one can provide a refere
 function resample_propagate!(
     ::Random.AbstractRNG,
     pc::ParticleContainer,
-    sampler::AbstractParticleSampler,
+    sampler::T,
     randcat=DEFAULT_RESAMPLER,
     ref::Union{Particle,Nothing}=nothing;
     weights=getweights(pc),
-)
+) where {T<:AbstractMCMC.AbstractSampler}
     # sample ancestor indices
     n = length(pc)
     nresamples = ref === nothing ? n : n - 1
@@ -233,11 +233,11 @@ end
 function resample_propagate!(
     rng::Random.AbstractRNG,
     pc::ParticleContainer,
-    sampler::AbstractParticleSampler,
+    sampler::T,
     resampler::ResampleWithESSThreshold,
     ref::Union{Particle,Nothing}=nothing;
     weights=getweights(pc),
-)
+) where {T<:AbstractMCMC.AbstractSampler}
     # Compute the effective sample size ``1 / ∑ wᵢ²`` with normalized weights ``wᵢ``
     ess = inv(sum(abs2, weights))