Use Distributions.loglikelihood instead of Distributions.logpdf (#153)

This PR replaces some occurrences of `logpdf` with `loglikelihood` since the former is intended to be used for a single sample only but we are interested in the log probability of individual and multiple samples.

JuliaStats/Distributions.jl#1144 allows us to use `loglikelihood` even for individual samples and arrays of samples (everything that can be sampled by `rand` should support the computation of `loglikelihood`). Currently, `logpdf` is misused (also in DistributionsAD) to compute arrays of log densities for multiple samples which are summed afterwards. Usually, this intermittent step can be avoided by summing the log densities directly (which is the default implementation in Distributions).

Similar issues and possible optimizations exist for `Bijectors.logpdf_with_trans` (see TuringLang/Bijectors.jl#120).

Co-authored-by: Hong Ge <[email protected]>

Author: devmotion

Project.toml

@@ -15,7 +15,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 AbstractMCMC = "1"
 Bijectors = "0.5.2, 0.6, 0.7, 0.8"
-Distributions = "0.22, 0.23"
+Distributions = "0.23.8"
 MacroTools = "0.5.1"
 NaturalSort = "1"
 ZygoteRules = "0.2"

src/context_implementations.jl

@@ -149,7 +149,7 @@ function observe(
     vi,
 )
     increment_num_produce!(vi)
-    return Distributions.logpdf(dist, value)
+    return Distributions.loglikelihood(dist, value)
 end
 
 # .~ functions
@@ -444,18 +444,33 @@ function dot_observe(
     vi,
 )
     increment_num_produce!(vi)
-    @debug "dot_observe" dist value
-    return sum(Distributions.logpdf(dist, value))
+    @debug "dist = $dist"
+    @debug "value = $value"
+    return Distributions.loglikelihood(dist, value)
 end
 function dot_observe(
     spl::Union{SampleFromPrior, SampleFromUniform},
-    dists::Union{Distribution, AbstractArray{<:Distribution}},
+    dists::Distribution,
     value::AbstractArray,
     vi,
 )
     increment_num_produce!(vi)
-    @debug "dot_observe" dists value
-    return sum(Distributions.logpdf.(dists, value))
+    @debug "dists = $dists"
+    @debug "value = $value"
+    return Distributions.loglikelihood(dists, value)
+end
+function dot_observe(
+    spl::Union{SampleFromPrior, SampleFromUniform},
+    dists::AbstractArray{<:Distribution},
+    value::AbstractArray,
+    vi,
+)
+    increment_num_produce!(vi)
+    @debug "dists = $dists"
+    @debug "value = $value"
+    return sum(zip(dists, value)) do (d, v)
+        Distributions.loglikelihood(d, v)
+    end
 end
 function dot_observe(
     spl::Sampler,