s/diffusion/diffusions

blackjax-devs · Nov 20, 2022 · d996dbe · d996dbe
1 parent fb02c8c
commit d996dbe
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 6 deletions.
diff --git a/blackjax/mcmc/diffusion.py → blackjax/mcmc/diffusions.py b/blackjax/mcmc/diffusion.py → blackjax/mcmc/diffusions.py
diff --git a/blackjax/mcmc/mala.py b/blackjax/mcmc/mala.py
@@ -5,7 +5,7 @@
 import jax
 import jax.numpy as jnp
 
-from blackjax.mcmc.diffusion import overdamped_langevin
+import blackjax.mcmc.diffusions as diffusions
 from blackjax.types import PRNGKey, PyTree
 
 __all__ = ["MALAState", "MALAInfo", "init", "kernel"]
@@ -83,7 +83,7 @@ def one_step(
 
         """
         grad_fn = jax.value_and_grad(logprob_fn)
-        integrator = overdamped_langevin(grad_fn)
+        integrator = diffusions.overdamped_langevin(grad_fn)
 
         key_integrator, key_rmh = jax.random.split(rng_key)
 

diff --git a/blackjax/sgmcmc/diffusion.py → blackjax/sgmcmc/diffusions.py b/blackjax/sgmcmc/diffusion.py → blackjax/sgmcmc/diffusions.py
diff --git a/blackjax/sgmcmc/sghmc.py b/blackjax/sgmcmc/sghmc.py
@@ -3,7 +3,7 @@
 
 import jax
 
-from blackjax.sgmcmc.diffusion import sghmc
+import blackjax.sgmcmc.diffusions as diffusions
 from blackjax.types import PRNGKey, PyTree
 from blackjax.util import generate_gaussian_noise
 
@@ -13,7 +13,7 @@
 def kernel(alpha: float = 0.01, beta: float = 0) -> Callable:
     """Stochastic gradient Hamiltonian Monte Carlo (SgHMC) algorithm."""
 
-    integrator = sghmc(alpha, beta)
+    integrator = diffusions.sghmc(alpha, beta)
 
     def one_step(
         rng_key: PRNGKey,

diff --git a/blackjax/sgmcmc/sgld.py b/blackjax/sgmcmc/sgld.py
@@ -1,7 +1,7 @@
 """Public API for the Stochastic gradient Langevin Dynamics kernel."""
 from typing import Callable
 
-from blackjax.sgmcmc.diffusion import overdamped_langevin
+import blackjax.sgmcmc.diffusions as diffusions
 from blackjax.types import PRNGKey, PyTree
 
 __all__ = ["kernel"]
@@ -10,7 +10,7 @@
 def kernel() -> Callable:
     """Stochastic gradient Langevin Dynamics (SgLD) algorithm."""
 
-    integrator = overdamped_langevin()
+    integrator = diffusions.overdamped_langevin()
 
     def one_step(
         rng_key: PRNGKey,