Repository: pyro-ppl/numpyro

Issue:
Support TFP mcmc kernels
TFP [kernels](https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/RandomWalkMetropolis#one_step) takes potential function in its construction and has the method [one_step](https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/RandomWalkMetropolis#one_step) which update `(current_state, previous_kernel_results, seed)`. This is equivalent to how [MCMCKernel](http://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.mcmc.MCMCKernel) works in NumPyro.

It should be straightforward to add wrappers in `contrib.tfp` for TFP MCMC kernels because both api are equivalent. However, in NumPyro, we support both `model` and `potential_fn` so I think the wrappers in `contrib.tfp.mcmc` should take the same name but `target_log_prob_fn` will be replaced by `model=None, potential_fn=None` (all other args, kwargs are kept the same).

While looking into this, I also think that we should make `HMCState`, `SAState` a default attribute of their corresponding classes. This centralizes all information related to kernel to 1 class, rather than separated into a pair of (state_definition, kernel_class) as in the current implementation.

WDYT about the above 2 issues @neerajprad?

After this, we can use tfp mcmc kernels for NumPyro models. :D


Discussion:

Resolving patch:
diff --git a/numpyro/contrib/tfp/mcmc.py b/numpyro/contrib/tfp/mcmc.py
new file mode 100644
--- /dev/null
+++ b/numpyro/contrib/tfp/mcmc.py
@@ -0,0 +1,229 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABCMeta
+from collections import namedtuple
+import inspect
+
+from jax import random, tree_map, vmap
+from jax.flatten_util import ravel_pytree
+import jax.numpy as jnp
+import tensorflow_probability.substrates.jax as tfp
+
+from numpyro.infer import init_to_uniform
+from numpyro.infer.mcmc import MCMCKernel
+from numpyro.infer.util import initialize_model
+from numpyro.util import identity
+
+TFPKernelState = namedtuple('TFPKernelState', ['z', 'kernel_results', 'rng_key'])
+
+
+def _extract_kernel_functions(kernel):
+
+    def init_fn(z, rng_key):
+        z_flat, _ = ravel_pytree(z)
+        results = kernel.bootstrap_results(z_flat)
+        return TFPKernelState(z, results, rng_key)
+
+    def sample_fn(state, model_args=(), model_kwargs=None):
+        rng_key, rng_key_transition = random.split(state.rng_key)
+        z_flat, unravel_fn = ravel_pytree(state.z)
+        z_new_flat, results = kernel.one_step(z_flat, state.kernel_results, seed=rng_key_transition)
+        return TFPKernelState(unravel_fn(z_new_flat), results, rng_key)
+
+    return init_fn, sample_fn
+
+
+def _make_log_prob_fn(potential_fn, unravel_fn):
+
+    def log_prob_fn(x):
+        # we deal with batched x in case the kernel is ReplicaExchangeMC
+        batch_shape = jnp.shape(x)[:-1]
+        if batch_shape:
+            flatten_result = vmap(lambda a: -potential_fn(unravel_fn(a)))(
+                jnp.reshape(x, (-1,) + jnp.shape(x)[-1:]))
+            return tree_map(lambda a: jnp.reshape(a, batch_shape + jnp.shape(a)[1:]),
+                            flatten_result)
+        else:
+            return - potential_fn(unravel_fn(x))
+
+    return log_prob_fn
+
+
+class _TFPKernelMeta(ABCMeta):
+    def __getitem__(cls, kernel_class):
+        assert issubclass(kernel_class, tfp.mcmc.TransitionKernel)
+        assert 'target_log_prob_fn' in inspect.getfullargspec(kernel_class).args, \
+            f"the first argument of {kernel_class} must be `target_log_prob_fn`"
+
+        _PyroKernel = type(kernel_class.__name__, (TFPKernel,), {})
+        _PyroKernel.kernel_class = kernel_class
+        return _PyroKernel
+
+
+class TFPKernel(MCMCKernel, metaclass=_TFPKernelMeta):
+    """
+    A thin wrapper for TensorFlow Probability (TFP) MCMC transition kernels.
+    The argument `target_log_prob_fn` in TFP is replaced by either `model`
+    or `potential_fn` (which is the negative of `target_log_prob_fn`).
+
+    This class can be used to convert a TFP kernel to a NumPyro-compatible one
+    as follows::
+
+        kernel = TFPKernel[tfp.mcmc.NoUTurnSampler](model, step_size=1.)
+
+    .. note:: By default, uncalibrated kernels will be inner kernels of the
+        :class:`~tensorflow_probability.substrates.jax.mcmc.MetropolisHastings` kernel.
+
+    .. note:: For :class:`~numpyro.contrib.tfp.mcmc.ReplicaExchangeMC`, TFP requires
+        that the shape of `step_size` of the inner kernel must be
+        `[len(inverse_temperatures), 1]` or `[len(inverse_temperatures), latent_size]`.
+
+    :param model: Python callable containing Pyro :mod:`~numpyro.primitives`.
+        If model is provided, `potential_fn` will be inferred using the model.
+    :param potential_fn: Python callable that computes the target potential energy
+        given input parameters. The input parameters to `potential_fn`
+        can be any python collection type, provided that `init_params` argument to
+        :meth:`init` has the same type.
+    :param callable init_strategy: a per-site initialization function.
+        See :ref:`init_strategy` section for available functions.
+    :param kernel_kwargs: other arguments to be passed to TFP kernel constructor.
+    """
+    kernel_class = None
+
+    def __init__(self, model=None, potential_fn=None, init_strategy=init_to_uniform,
+                 **kernel_kwargs):
+        if not (model is None) ^ (potential_fn is None):
+            raise ValueError('Only one of `model` or `potential_fn` must be specified.')
+        self._model = model
+        self._potential_fn = potential_fn
+        self._kernel_kwargs = kernel_kwargs
+        self._init_strategy = init_strategy
+        # Set on first call to init
+        self._init_fn = None
+        self._postprocess_fn = None
+        self._sample_fn = None
+
+    def _init_state(self, rng_key, model_args, model_kwargs, init_params):
+        if self._model is not None:
+            init_params, potential_fn, postprocess_fn, model_trace = initialize_model(
+                rng_key,
+                self._model,
+                init_strategy=self._init_strategy,
+                dynamic_args=True,
+                model_args=model_args,
+                model_kwargs=model_kwargs)
+            init_params = init_params.z
+            if self._init_fn is None:
+                _, unravel_fn = ravel_pytree(init_params)
+                kernel = self.kernel_class(
+                    _make_log_prob_fn(potential_fn(*model_args, **model_kwargs), unravel_fn),
+                    **self._kernel_kwargs)
+                # Uncalibrated... kernels have to used inside MetropolisHastings, see
+                # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/UncalibratedLangevin
+                if self.kernel_class.__name__.startswith("Uncalibrated"):
+                    kernel = tfp.mcmc.MetropolisHastings(kernel)
+                self._init_fn, self._sample_fn = _extract_kernel_functions(kernel)
+            self._postprocess_fn = postprocess_fn
+        elif self._init_fn is None:
+            _, unravel_fn = ravel_pytree(init_params)
+            kernel = self.kernel_class(
+                _make_log_prob_fn(self._potential_fn, unravel_fn),
+                **self._kernel_kwargs)
+            if self.kernel_class.__name__.startswith("Uncalibrated"):
+                kernel = tfp.mcmc.MetropolisHastings(kernel)
+            self._init_fn, self._sample_fn = _extract_kernel_functions(kernel)
+        return init_params
+
+    @property
+    def model(self):
+        return self._model
+
+    @property
+    def sample_field(self):
+        return 'z'
+
+    @property
+    def default_fields(self):
+        return ('z',)
+
+    def get_diagnostics_str(self, state):
+        """
+        Given the current `state`, returns the diagnostics string to
+        be added to progress bar for diagnostics purpose.
+        """
+        return ''
+
+    def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={}):
+        # non-vectorized
+        if rng_key.ndim == 1:
+            rng_key, rng_key_init_model = random.split(rng_key)
+        # vectorized
+        else:
+            rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1)
+        init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
+        if self._potential_fn and init_params is None:
+            raise ValueError('Valid value of `init_params` must be provided with'
+                             ' `target_log_prob_fn`.')
+
+        if rng_key.ndim == 1:
+            init_state = self._init_fn(init_params, rng_key)
+        else:
+            # XXX it is safe to run hmc_init_fn under vmap despite that hmc_init_fn changes some
+            # nonlocal variables: momentum_generator, wa_update, trajectory_len, max_treedepth,
+            # wa_steps because those variables do not depend on traced args: init_params, rng_key.
+            init_state = vmap(self._init_fn)(init_params, rng_key)
+            sample_fn = vmap(self._sample_fn, in_axes=(0, None, None))
+            self._sample_fn = sample_fn
+        return init_state
+
+    def postprocess_fn(self, args, kwargs):
+        if self._postprocess_fn is None:
+            return identity
+        return self._postprocess_fn(*args, **kwargs)
+
+    def sample(self, state, model_args, model_kwargs):
+        """
+        Run the kernel from the given :data:`~numpyro.contrib.tfp.mcmc.TFPKernelState`
+        and return the resulting :data:`~numpyro.contrib.tfp.mcmc.TFPKernelState`.
+
+        :param TFPKernelState state: Represents the current state.
+        :param model_args: Arguments provided to the model.
+        :param model_kwargs: Keyword arguments provided to the model.
+        :return: Next `state` after running the kernel.
+        """
+        return self._sample_fn(state, model_args, model_kwargs)
+
+
+__all__ = ['TFPKernel']
+for _name, _Kernel in tfp.mcmc.__dict__.items():
+    if not isinstance(_Kernel, type):
+        continue
+    if not issubclass(_Kernel, tfp.mcmc.TransitionKernel):
+        continue
+    if 'target_log_prob_fn' not in inspect.getfullargspec(_Kernel).args:
+        continue
+
+    _PyroKernel = TFPKernel[_Kernel]
+    _PyroKernel.__module__ = __name__
+    locals()[_name] = _PyroKernel
+
+    _PyroKernel.__doc__ = '''
+    Wraps `{}.{} <https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/{}>`_
+    with :class:`~numpyro.contrib.tfp.mcmc.TFPKernel`. The first argument `target_log_prob_fn`
+    in TFP kernel construction is replaced by either `model` or `potential_fn`.
+    '''.format(_Kernel.__module__, _Kernel.__name__, _Kernel.__name__)
+
+    __all__.append(_name)
+
+
+# Create sphinx documentation.
+__doc__ = '\n\n'.join([
+
+    '''
+    {0}
+    ----------------------------------------------------------------
+    .. autoclass:: numpyro.contrib.tfp.mcmc.{0}
+    '''.format(_name)
+    for _name in __all__[:1] + sorted(__all__)
+])
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -1,7 +1,6 @@
 from collections import namedtuple
 import math
 import os
-import warnings
 
 from jax import device_put, lax, partial, random, vmap
 from jax.dtypes import canonicalize_dtype
@@ -348,7 +347,7 @@ class HMC(MCMCKernel):
     :param potential_fn: Python callable that computes the potential energy
         given input parameters. The input parameters to `potential_fn` can be
         any python collection type, provided that `init_params` argument to
-        `init_kernel` has the same type.
+        :meth:`init` has the same type.
     :param kinetic_fn: Python callable that returns the kinetic energy given
         inverse mass matrix and momentum. If not provided, the default is
         euclidean kinetic energy.
@@ -412,11 +411,6 @@ def _init_state(self, rng_key, model_args, model_kwargs, init_params):
                 init_strategy=self._init_strategy,
                 model_args=model_args,
                 model_kwargs=model_kwargs)
-            if any(v['type'] == 'param' for v in model_trace.values()):
-                warnings.warn("'param' sites will be treated as constants during inference. To define "
-                              "an improper variable, please use a 'sample' site with log probability "
-                              "masked out. For example, `sample('x', dist.LogNormal(0, 1).mask(False)` "
-                              "means that `x` has improper distribution over the positive domain.")
             if self._init_fn is None:
                 self._init_fn, self._sample_fn = hmc(potential_fn_gen=potential_fn,
                                                      kinetic_fn=self._kinetic_fn,
diff --git a/numpyro/infer/sa.py b/numpyro/infer/sa.py
--- a/numpyro/infer/sa.py
+++ b/numpyro/infer/sa.py
@@ -225,7 +225,7 @@ class SA(MCMCKernel):
     :param potential_fn: Python callable that computes the potential energy
         given input parameters. The input parameters to `potential_fn` can be
         any python collection type, provided that `init_params` argument to
-        `init_kernel` has the same type.
+        :meth:`init` has the same type.
     :param int adapt_state_size: The number of points to generate proposal
         distribution. Defaults to 2 times latent size.
     :param bool dense_mass:  A flag to decide if mass matrix is dense or


Test patch:
diff --git a/test/contrib/test_tfp.py b/test/contrib/test_tfp.py
--- a/test/contrib/test_tfp.py
+++ b/test/contrib/test_tfp.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import inspect
+import os
 
 from numpy.testing import assert_allclose
 import pytest
@@ -12,6 +13,7 @@
 import numpyro
 import numpyro.distributions as dist
 from numpyro.infer import MCMC, NUTS
+from numpyro.infer.reparam import TransformReparam
 
 
 # XXX: for some reasons, pytest raises ImportWarning when we import tfp
@@ -103,3 +105,81 @@ def model(data):
     mcmc.print_summary()
     samples = mcmc.get_samples()
     assert_allclose(jnp.mean(samples['p_latent'], 0), true_probs, atol=0.05)
+
+
+def make_kernel_fn(target_log_prob_fn):
+    import tensorflow_probability.substrates.jax as tfp
+
+    return tfp.mcmc.HamiltonianMonteCarlo(
+        target_log_prob_fn=target_log_prob_fn,
+        step_size=0.5 / jnp.sqrt(0.5 ** jnp.arange(4)[..., None]), num_leapfrog_steps=5)
+
+
+@pytest.mark.parametrize('kernel, kwargs', [
+    ('HamiltonianMonteCarlo', dict(step_size=0.05, num_leapfrog_steps=10)),
+    ('NoUTurnSampler', dict(step_size=0.05)),
+    ('RandomWalkMetropolis', dict()),
+    ('SliceSampler', dict(step_size=1.0, max_doublings=5)),
+    ('UncalibratedHamiltonianMonteCarlo', dict(step_size=0.05, num_leapfrog_steps=10)),
+    ('UncalibratedRandomWalk', dict()),
+])
+@pytest.mark.filterwarnings("ignore:can't resolve package")
+# TODO: remove after https://github.com/tensorflow/probability/issues/1072 is resolved
+@pytest.mark.filterwarnings("ignore:Explicitly requested dtype")
+def test_mcmc_kernels(kernel, kwargs):
+    from numpyro.contrib.tfp import mcmc
+    kernel_class = getattr(mcmc, kernel)
+
+    true_coef = 0.9
+    num_warmup, num_samples = 1000, 1000
+
+    def model(data):
+        alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
+        with numpyro.handlers.reparam(config={'loc': TransformReparam()}):
+            loc = numpyro.sample('loc', dist.Uniform(0, alpha))
+        numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
+
+    data = true_coef + random.normal(random.PRNGKey(0), (1000,))
+    tfp_kernel = kernel_class(model=model, **kwargs)
+    mcmc = MCMC(tfp_kernel, num_warmup=num_warmup, num_samples=num_samples)
+    mcmc.warmup(random.PRNGKey(2), data, collect_warmup=True)
+    warmup_samples = mcmc.get_samples()
+    mcmc.run(random.PRNGKey(3), data)
+    samples = mcmc.get_samples()
+    assert len(warmup_samples['loc']) == num_warmup
+    assert len(samples['loc']) == num_samples
+    assert_allclose(jnp.mean(samples['loc'], 0), true_coef, atol=0.05)
+
+
+@pytest.mark.parametrize('kernel, kwargs', [
+    ('MetropolisAdjustedLangevinAlgorithm', dict(step_size=1.0)),
+    ('RandomWalkMetropolis', dict()),
+    ('SliceSampler', dict(step_size=1.0, max_doublings=5)),
+    ('UncalibratedLangevin', dict(step_size=0.1)),
+    ('ReplicaExchangeMC', dict(inverse_temperatures=0.5 ** jnp.arange(4), make_kernel_fn=make_kernel_fn))
+])
+@pytest.mark.parametrize('num_chains', [1, 2])
+@pytest.mark.skipif('XLA_FLAGS' not in os.environ, reason='without this mark, we have duplicated tests in Travis')
+@pytest.mark.filterwarnings("ignore:There are not enough devices:UserWarning")
+@pytest.mark.filterwarnings("ignore:can't resolve package")
+# TODO: remove after https://github.com/tensorflow/probability/issues/1072 is resolved
+@pytest.mark.filterwarnings("ignore:Explicitly requested dtype")
+def test_unnormalized_normal_chain(kernel, kwargs, num_chains):
+    from numpyro.contrib.tfp import mcmc
+
+    kernel_class = getattr(mcmc, kernel)
+
+    true_mean, true_std = 1., 0.5
+    warmup_steps, num_samples = (1000, 8000)
+
+    def potential_fn(z):
+        return 0.5 * ((z - true_mean) / true_std) ** 2
+
+    init_params = jnp.array(0.) if num_chains == 1 else jnp.array([0., 2.])
+    tfp_kernel = kernel_class(potential_fn=potential_fn, **kwargs)
+    mcmc = MCMC(tfp_kernel, warmup_steps, num_samples, num_chains=num_chains, progress_bar=False)
+    mcmc.run(random.PRNGKey(0), init_params=init_params)
+    mcmc.print_summary()
+    hmc_states = mcmc.get_samples()
+    assert_allclose(jnp.mean(hmc_states), true_mean, rtol=0.07)
+    assert_allclose(jnp.std(hmc_states), true_std, rtol=0.07)