Add logjac to logdensity_fn

blackjax-devs · Oct 30, 2024 · 8611f8f · 8611f8f
1 parent b107f9f
commit 8611f8f
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/docs/examples/howto_sample_multiple_chains.md b/docs/examples/howto_sample_multiple_chains.md
@@ -57,8 +57,9 @@ observed = np.random.normal(loc, scale, size=1_000)
 def logdensity_fn(loc, log_scale, observed=observed):
     """Univariate Normal"""
     scale = jnp.exp(log_scale)
+    logjac = log_scale
     logpdf = stats.norm.logpdf(observed, loc, scale)
-    return jnp.sum(logpdf)
+    return logjac + jnp.sum(logpdf)
 
 
 def logdensity(x):

diff --git a/docs/examples/quickstart.md b/docs/examples/quickstart.md
@@ -48,8 +48,9 @@ observed = np.random.normal(loc, scale, size=1_000)
 def logdensity_fn(loc, log_scale, observed=observed):
     """Univariate Normal"""
     scale = jnp.exp(log_scale)
+    logjac = log_scale
     logpdf = stats.norm.logpdf(observed, loc, scale)
-    return jnp.sum(logpdf)
+    return logjac + jnp.sum(logpdf)
 
 
 logdensity = lambda x: logdensity_fn(**x)

diff --git a/tests/smc/__init__.py b/tests/smc/__init__.py
@@ -9,7 +9,7 @@ def logdensity_by_observation(self, log_scale, coefs, preds, x):
         scale = jnp.exp(log_scale)
         y = jnp.dot(x, coefs)
         logpdf = stats.norm.logpdf(preds, y, scale)
-        return logpdf
+        return log_scale + logpdf
 
     def logdensity_fn(self, log_scale, coefs, preds, x):
         """Linear regression"""

diff --git a/tests/smc/test_inner_kernel_tuning.py b/tests/smc/test_inner_kernel_tuning.py
@@ -69,7 +69,7 @@ def logdensity_fn(self, log_scale, coefs, preds, x):
         scale = jnp.exp(log_scale)
         y = jnp.dot(x, coefs)
         logpdf = stats.norm.logpdf(preds, y, scale)
-        return jnp.sum(logpdf)
+        return log_scale + jnp.sum(logpdf)
 
     def test_smc_inner_kernel_adaptive_tempered(self):
         self.smc_inner_kernel_tuning_test_case(