A draft PR on the removal of Algorithm classes. variant 2

blackjax/__init__.py

@@ -11,7 +11,7 @@
 from .mcmc.dynamic_hmc import dynamic_hmc
 from .mcmc.elliptical_slice import elliptical_slice
 from .mcmc.ghmc import ghmc
-from .mcmc.hmc import hmc
+from .mcmc import hmc
 from .mcmc.mala import mala
 from .mcmc.marginal_latent_gaussian import mgrad_gaussian
 from .mcmc.mclmc import mclmc

blackjax/adaptation/pathfinder_adaptation.py

@@ -138,7 +138,7 @@ def final(warmup_state: PathfinderAdaptationState) -> tuple[float, Array]:
 
 
 def pathfinder_adaptation(
-    algorithm: Union[mcmc.hmc.hmc, mcmc.nuts.nuts],
+    algorithm,
     logdensity_fn: Callable,
     initial_step_size: float = 1.0,
     target_acceptance_rate: float = 0.80,

blackjax/adaptation/window_adaptation.py

@@ -243,7 +243,7 @@ def final(warmup_state: WindowAdaptationState) -> tuple[float, Array]:
 
 
 def window_adaptation(
-    algorithm: Union[mcmc.hmc.hmc, mcmc.nuts.nuts],
+    algorithm,
     logdensity_fn: Callable,
     is_mass_matrix_diagonal: bool = True,
     initial_step_size: float = 1.0,

blackjax/mcmc/hmc.py

@@ -29,7 +29,7 @@
     "HMCInfo",
     "init",
     "build_kernel",
-    "hmc",
+    "as_sampling_algorithm",
 ]
 
 
@@ -150,7 +150,14 @@ def kernel(
     return kernel
 
 
-class hmc:
+def as_sampling_algorithm(logdensity_fn: Callable,
+        step_size: float,
+        inverse_mass_matrix: metrics.MetricTypes,
+        num_integration_steps: int,
+        *,
+        divergence_threshold: int = 1000,
+        integrator: Callable = integrators.velocity_verlet,
+    ) -> SamplingAlgorithm:
     """Implements the (basic) user interface for the HMC kernel.
 
     The general hmc kernel builder (:meth:`blackjax.mcmc.hmc.build_kernel`, alias
@@ -225,36 +232,24 @@ class hmc:
     A ``SamplingAlgorithm``.
     """
 
-    init = staticmethod(init)
-    build_kernel = staticmethod(build_kernel)
 
-    def __new__(  # type: ignore[misc]
-        cls,
-        logdensity_fn: Callable,
-        step_size: float,
-        inverse_mass_matrix: metrics.MetricTypes,
-        num_integration_steps: int,
-        *,
-        divergence_threshold: int = 1000,
-        integrator: Callable = integrators.velocity_verlet,
-    ) -> SamplingAlgorithm:
-        kernel = cls.build_kernel(integrator, divergence_threshold)
-
-        def init_fn(position: ArrayLikeTree, rng_key=None):
-            del rng_key
-            return cls.init(position, logdensity_fn)
-
-        def step_fn(rng_key: PRNGKey, state):
-            return kernel(
-                rng_key,
-                state,
-                logdensity_fn,
-                step_size,
-                inverse_mass_matrix,
-                num_integration_steps,
-            )
-
-        return SamplingAlgorithm(init_fn, step_fn)
+    kernel = build_kernel(integrator, divergence_threshold)
+
+    def init_fn(position: ArrayLikeTree, rng_key=None):
+        del rng_key
+        return init(position, logdensity_fn)
+
+    def step_fn(rng_key: PRNGKey, state):
+        return kernel(
+            rng_key,
+            state,
+            logdensity_fn,
+            step_size,
+            inverse_mass_matrix,
+            num_integration_steps,
+        )
+
+    return SamplingAlgorithm(init_fn, step_fn)
 
 
 def hmc_proposal(

tests/mcmc/test_sampling.py

@@ -41,7 +41,7 @@ def rmh_proposal_distribution(rng_key, position):
 
 regression_test_cases = [
     {
-        "algorithm": blackjax.hmc,
+        "algorithm": blackjax.hmc.as_sampling_algorithm,
         "initial_position": {"log_scale": 0.0, "coefs": 4.0},
         "parameters": {"num_integration_steps": 90},
         "num_warmup_steps": 1_000,

tests/test_benchmarks.py

@@ -36,7 +36,6 @@ def run_regression(algorithm, **parameters):
     logdensity_fn = lambda x: logdensity_fn_(**x)
 
     warmup_key, inference_key = jax.random.split(rng_key, 2)
-
     warmup = blackjax.window_adaptation(
         algorithm,
         logdensity_fn,
@@ -46,7 +45,7 @@ def run_regression(algorithm, **parameters):
     (state, parameters), _ = warmup.run(
         warmup_key, {"log_scale": 0.0, "coefs": 2.0}, 1000
     )
-    inference_algorithm = algorithm(logdensity_fn, **parameters)
+    inference_algorithm = algorithm.as_sampling_algorithm(logdensity_fn, **parameters)
 
     _, states, _ = run_inference_algorithm(
         inference_key, state, inference_algorithm, 10_000

tests/test_compilation.py

@@ -32,7 +32,7 @@ def logdensity_fn(x):
         rng_key = jax.random.key(0)
         state = blackjax.hmc.init(1.0, logdensity_fn)
 
-        kernel = blackjax.hmc(
+        kernel = blackjax.hmc.as_sampling_algorithm(
             logdensity_fn,
             step_size=1e-2,
             inverse_mass_matrix=jnp.array([1.0]),
@@ -92,7 +92,7 @@ def logdensity_fn(x):
             num_integration_steps=10,
         )
         (state, parameters), _ = warmup.run(rng_key, 1.0, num_steps=100)
-        kernel = jax.jit(blackjax.hmc(logdensity_fn, **parameters).step)
+        kernel = jax.jit(blackjax.hmc.as_sampling_algorithm(logdensity_fn, **parameters).step)
 
         for _ in range(10):
             rng_key, sample_key = jax.random.split(rng_key)

tests/test_util.py

@@ -3,15 +3,15 @@
 import jax.numpy as jnp
 from absl.testing import absltest, parameterized
 
-from blackjax.mcmc.hmc import hmc
+from blackjax.mcmc.hmc import as_sampling_algorithm
 from blackjax.util import run_inference_algorithm
 
 
 class RunInferenceAlgorithmTest(chex.TestCase):
     def setUp(self):
         super().setUp()
         self.key = jax.random.key(42)
-        self.algorithm = hmc(
+        self.algorithm = as_sampling_algorithm(
             logdensity_fn=self.logdensity_fn,
             inverse_mass_matrix=jnp.eye(2),
             step_size=1.0,