[Draft] Nested sampling implementation

blackjax/__init__.py

@@ -22,12 +22,14 @@
 from .mcmc import nuts as _nuts
 from .mcmc import periodic_orbital, random_walk
 from .mcmc import rmhmc as _rmhmc
+from .mcmc import ss
 from .mcmc.random_walk import additive_step_random_walk as _additive_step_random_walk
 from .mcmc.random_walk import (
     irmh_as_top_level_api,
     normal_random_walk,
     rmh_as_top_level_api,
 )
+from .ns import nss as _nss
 from .optimizers import dual_averaging, lbfgs
 from .sgmcmc import csgld as _csgld
 from .sgmcmc import sghmc as _sghmc
@@ -109,6 +111,8 @@ def generate_top_level_api_from(module):
 
 additive_step_random_walk.register_factory("normal_random_walk", normal_random_walk)
 
+hrss = GenerateSamplingAPI(ss.hrss_as_top_level_api, ss.init, ss.build_kernel)
+
 mclmc = generate_top_level_api_from(_mclmc)
 elliptical_slice = generate_top_level_api_from(_elliptical_slice)
 ghmc = generate_top_level_api_from(_ghmc)
@@ -122,6 +126,8 @@ def generate_top_level_api_from(module):
 inner_kernel_tuning = generate_top_level_api_from(_inner_kernel_tuning)
 partial_posteriors_smc = generate_top_level_api_from(_partial_posteriors_smc)
 
+nss = generate_top_level_api_from(_nss)
+
 smc_family = [tempered_smc, adaptive_tempered_smc, partial_posteriors_smc]
 "Step_fn returning state has a .particles attribute"