Improve integration with external samplers

numpyro/_typing.py

@@ -35,4 +35,13 @@
 """A generic type for a pytree, i.e. a nested structure of lists, tuples, dicts, and arrays."""
 
 
+PositionDict: TypeAlias = dict[str, jax.Array]
+"""An unconstrained position dict keyed by sample-site name.
+
+Used as the canonical input/output type for log-density and postprocess
+callables exposed to external samplers (see
+:class:`~numpyro.infer.LogDensityInfo` and
+:class:`~numpyro.infer.ExternalKernel`)."""
+
+
 NumLikeT = TypeVar("NumLikeT", bound=NumLike)

numpyro/infer/__init__.py

@@ -12,6 +12,7 @@
     TraceMeanField_ELBO,
 )
 from numpyro.infer.ensemble import AIES, ESS
+from numpyro.infer.external import ExternalKernel
 from numpyro.infer.hmc import HMC, NUTS
 from numpyro.infer.hmc_gibbs import HMCECS, DiscreteHMCGibbs, HMCGibbs
 from numpyro.infer.importance import psis_diagnostic
@@ -27,30 +28,42 @@
 from numpyro.infer.mixed_hmc import MixedHMC
 from numpyro.infer.sa import SA
 from numpyro.infer.svi import SVI
-from numpyro.infer.util import Predictive, log_likelihood
+from numpyro.infer.util import (
+    LogDensityInfo,
+    Predictive,
+    constrain_samples,
+    get_log_density_fn,
+    initialize_model,
+    log_likelihood,
+)
 
 from . import autoguide, calibration, reparam
 
 __all__ = [
     "AIES",
     "autoguide",
     "calibration",
+    "constrain_samples",
+    "get_log_density_fn",
     "init_to_feasible",
     "init_to_mean",
     "init_to_median",
     "init_to_sample",
     "init_to_uniform",
     "init_to_value",
+    "initialize_model",
     "log_likelihood",
     "psis_diagnostic",
     "reparam",
     "BarkerMH",
     "DiscreteHMCGibbs",
     "ELBO",
     "ESS",
+    "ExternalKernel",
     "HMC",
     "HMCECS",
     "HMCGibbs",
+    "LogDensityInfo",
     "MCMC",
     "MixedHMC",
     "NUTS",