Repository: pyro-ppl/numpyro

Issue:
Using TFP distributions directly in `sample` primitive
Quoted from [the forum](https://forum.pyro.ai/t/sample-from-the-mixture-same-family-distribution/3178/6): The current wrappers in `numpyro.contrib.tfp` seems to not work for special distributions like MixtureSameFamily. For now, to make it work, we'll need to do
```python
from numpyro.contrib.tfp.distributions import TFPDistribution
import tensorflow_probability...distributions as tfd

def model():
    x = sample("x", TFPDistribution(tfd.MixtureSampleFamily(...)))
```
In the long term, I think it would be nice to use TFP distributions directly in the `sample` primitive:
```python
import tensorflow_probability...distributions as tfd

def model():
    x = sample("x", tfd.MixtureSampleFamily(...))
```

I believe it is enough to modify [`sample` primitive](https://github.com/pyro-ppl/numpyro/blob/e55f0d41c9eba48a10e88fb403a5e016f18857e6/numpyro/primitives.py#L104) such that if `fn` is not a NumPyro distribution, we will check if `tensorflow_probability` is available, then trying to check if `fn` is a TFP distribution. If yes, then we replace `fn` by `numpyro.contrib.tfp.distributions.TFPDistribution[type(fn)](**fn.parameters)` or by
```python
new_fn = TFPDistribution()
new_fn.tfp_dist = fn
new_fn.tfd_class = type(fn)  # and extra repr if needed
```


Discussion:

Resolving patch:
diff --git a/numpyro/contrib/tfp/distributions.py b/numpyro/contrib/tfp/distributions.py
--- a/numpyro/contrib/tfp/distributions.py
+++ b/numpyro/contrib/tfp/distributions.py
@@ -1,10 +1,12 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-import inspect
+from functools import lru_cache
+import warnings
 
 import numpy as np
 
+import jax
 import jax.numpy as jnp
 from tensorflow_probability.substrates.jax import bijectors as tfb, distributions as tfd
 
@@ -105,18 +107,68 @@ def _transform_to_bijector_constraint(constraint):
     return BijectorTransform(constraint.bijector)
 
 
+def _onehot_enumerate_support(self, expand=True):
+    n = self.event_shape[-1]
+    values = jnp.identity(n, dtype=jnp.result_type(self.dtype))
+    values = values.reshape((n,) + (1,) * len(self.batch_shape) + (n,))
+    if expand:
+        values = jnp.broadcast_to(values, (n,) + self.batch_shape + (n,))
+    return values
+
+
 class _TFPDistributionMeta(type(NumPyroDistribution)):
+    @lru_cache(maxsize=None)
     def __getitem__(cls, tfd_class):
         assert issubclass(tfd_class, tfd.Distribution)
 
+        tfd_class_name = tfd_class.__name__
+
         def init(self, *args, **kwargs):
+            warnings.warn(
+                "Importing distributions from numpyro.contrib.tfp.distributions is "
+                "deprecated. You should import distributions directly from "
+                "tensorflow_probability.substrates.jax.distributions instead.",
+                FutureWarning,
+            )
             self.tfp_dist = tfd_class(*args, **kwargs)
 
-        init.__signature__ = inspect.signature(tfd_class.__init__)
-
-        _PyroDist = type(tfd_class.__name__, (TFPDistribution,), {})
-        _PyroDist.tfd_class = tfd_class
+        _PyroDist = type(tfd_class_name, (TFPDistribution,), {})
         _PyroDist.__init__ = init
+
+        if tfd_class is tfd.InverseGamma:
+            _PyroDist.arg_constraints = {
+                "concentration": constraints.positive,
+                "scale": constraints.positive,
+            }
+        elif tfd_class is tfd.OneHotCategorical:
+            _PyroDist.arg_constraints = {"logits": constraints.real_vector}
+            _PyroDist.has_enumerate_support = True
+            _PyroDist.support = constraints.simplex
+            _PyroDist.is_discrete = True
+            _PyroDist.enumerate_support = _onehot_enumerate_support
+        elif tfd_class is tfd.OrderedLogistic:
+            _PyroDist.arg_constraints = {
+                "cutpoints": constraints.ordered_vector,
+                "loc": constraints.real,
+            }
+        elif tfd_class is tfd.Pareto:
+            _PyroDist.arg_constraints = {
+                "concentration": constraints.positive,
+                "scale": constraints.positive,
+            }
+        else:
+            if hasattr(numpyro_dist, tfd_class_name):
+                numpyro_dist_class = getattr(numpyro_dist, tfd_class_name)
+                # resolve FooProbs/FooLogits namespaces
+                numpyro_dist_class = getattr(
+                    numpyro_dist, f"{tfd_class_name}Logits", numpyro_dist_class
+                )
+                _PyroDist.arg_constraints = numpyro_dist_class.arg_constraints
+                _PyroDist.has_enumerate_support = (
+                    numpyro_dist_class.has_enumerate_support
+                )
+                _PyroDist.enumerate_support = numpyro_dist_class.enumerate_support
+
         return _PyroDist
 
 
@@ -132,8 +184,6 @@ class TFPDistribution(NumPyroDistribution, metaclass=_TFPDistributionMeta):
 
     """
 
-    tfd_class = None
-
     def __getattr__(self, name):
         # return parameters from the constructor
         if name in self.tfp_dist.parameters:
@@ -187,41 +237,15 @@ def is_discrete(self):
         # XXX: this should cover most cases
         return self.support is None
 
+    def tree_flatten(self):
+        return jax.tree_util.tree_flatten(self.tfp_dist)
 
-InverseGamma = TFPDistribution[tfd.InverseGamma]
-InverseGamma.arg_constraints = {
-    "concentration": constraints.positive,
-    "scale": constraints.positive,
-}
-
-
-def _onehot_enumerate_support(self, expand=True):
-    n = self.event_shape[-1]
-    values = jnp.identity(n, dtype=jnp.result_type(self.dtype))
-    values = values.reshape((n,) + (1,) * len(self.batch_shape) + (n,))
-    if expand:
-        values = jnp.broadcast_to(values, (n,) + self.batch_shape + (n,))
-    return values
-
-
-OneHotCategorical = TFPDistribution[tfd.OneHotCategorical]
-OneHotCategorical.arg_constraints = {"logits": constraints.real_vector}
-OneHotCategorical.has_enumerate_support = True
-OneHotCategorical.support = constraints.simplex
-OneHotCategorical.is_discrete = True
-OneHotCategorical.enumerate_support = _onehot_enumerate_support
-
-OrderedLogistic = TFPDistribution[tfd.OrderedLogistic]
-OrderedLogistic.arg_constraints = {
-    "cutpoints": constraints.ordered_vector,
-    "loc": constraints.real,
-}
-
-Pareto = TFPDistribution[tfd.Pareto]
-Pareto.arg_constraints = {
-    "concentration": constraints.positive,
-    "scale": constraints.positive,
-}
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        fn = jax.tree_util.tree_unflatten(aux_data, params)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=FutureWarning)
+            return TFPDistribution[fn.__class__](**fn.parameters)
 
 
 __all__ = ["BijectorConstraint", "BijectorTransform", "TFPDistribution"]
@@ -234,22 +258,9 @@ def _onehot_enumerate_support(self, expand=True):
     if _Dist is tfd.Distribution:
         continue
 
-    try:
-        _PyroDist = locals()[_name]
-    except KeyError:
-        _PyroDist = TFPDistribution[_Dist]
-        _PyroDist.__module__ = __name__
-        if hasattr(numpyro_dist, _name):
-            numpyro_dist_class = getattr(numpyro_dist, _name)
-            # resolve FooProbs/FooLogits namespaces
-            if type(numpyro_dist_class).__name__ == "function":
-                if not hasattr(numpyro_dist, _name + "Logits"):
-                    continue
-                numpyro_dist_class = getattr(numpyro_dist, _name + "Logits")
-            _PyroDist.arg_constraints = numpyro_dist_class.arg_constraints
-            _PyroDist.has_enumerate_support = numpyro_dist_class.has_enumerate_support
-            _PyroDist.enumerate_support = numpyro_dist_class.enumerate_support
-        locals()[_name] = _PyroDist
+    _PyroDist = TFPDistribution[_Dist]
+    _PyroDist.__module__ = __name__
+    locals()[_name] = _PyroDist
 
     _PyroDist.__doc__ = """
     Wraps `{}.{} <https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/distributions/{}>`_
@@ -271,6 +282,6 @@ def _onehot_enumerate_support(self, expand=True):
     """.format(
             _name
         )
-        for _name in __all__[:_len_all] + sorted(__all__[_len_all:])
+        for _name in __all__[:_len_all]
     ]
 )
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -131,6 +131,44 @@ def sample(
         + "_unobserved"`` which should be used by guides.
     :return: sample from the stochastic `fn`.
     """
+    if not isinstance(fn, numpyro.distributions.Distribution):
+        type_error = TypeError(
+            "It looks like you tried to use a fn that isn't an instance of "
+            "numpyro.distributions.Distribution, funsor.Funsor or "
+            "tensorflow_probability.distributions.Distribution. If you're using "
+            "funsor or tensorflow_probability, make sure they are correctly installed."
+        )
+
+        # fn can be a funsor.Funsor, but this won't be installed for all users
+        try:
+            from funsor import Funsor
+        except ImportError:
+            Funsor = None
+
+        # if Funsor import failed, or fn is not a Funsor it's also possible fn could be
+        # a tensorflow_probability distribution
+        if Funsor is None or not isinstance(fn, Funsor):
+            try:
+                from tensorflow_probability.substrates.jax import distributions as tfd
+
+                from numpyro.contrib.tfp.distributions import TFPDistribution
+            except ImportError:
+                # if tensorflow_probability fails to import here, then fn is not a
+                # numpyro Distribution or a Funsor, and it can't have been a tfp
+                # distribution either, so raising TypeError is ok
+                raise type_error
+
+            if isinstance(fn, tfd.Distribution):
+                with warnings.catch_warnings():
+                    # ignore FutureWarnings when instantiating TFPDistribution
+                    warnings.simplefilter("ignore", category=FutureWarning)
+                    # if fn is a tfp distribution we need to wrap it
+                    fn = TFPDistribution[fn.__class__](**fn.parameters)
+            else:
+                # if tensorflow_probability imported, but fn is not tfd.Distribution we
+                # still need to raise a type error
+                raise type_error
+
     # if there are no active Messengers, we just draw a sample and return it as expected:
     if not _PYRO_STACK:
         return fn(rng_key=rng_key, sample_shape=sample_shape)


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
@@ -1,7 +1,6 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-import inspect
 import os
 
 from numpy.testing import assert_allclose
@@ -17,71 +16,49 @@
 from numpyro.infer.reparam import TransformReparam
 
 
-# XXX: for some reasons, pytest raises ImportWarning when we import tfp
-@pytest.mark.filterwarnings("ignore:can't resolve package")
-def test_api_consistent():
-    from numpyro.contrib.tfp import distributions as tfd
-
-    for name in tfd.__all__:
-        if name in numpyro.distributions.__all__:
-            tfp_dist = getattr(tfd, name)
-            numpyro_dist = getattr(numpyro.distributions, name)
-            if type(numpyro_dist).__name__ == "function":
-                numpyro_dist = getattr(numpyro.distributions, name + "Logits")
-            for p in tfp_dist.arg_constraints:
-                assert p in inspect.getfullargspec(tfp_dist.__init__)[0]
-
-
 def test_dist_pytree():
-    from numpyro.contrib.tfp import distributions as tfd
+    from tensorflow_probability.substrates.jax import distributions as tfd
 
-    def f(x):
-        return tfd.Normal(x, 1)
-
-    jit(f)(0)
+    from numpyro.contrib.tfp.distributions import TFPDistribution
 
+    @jit
+    def f(x):
+        with numpyro.handlers.seed(rng_seed=0), numpyro.handlers.trace() as tr:
+            numpyro.sample("x", tfd.Normal(x, 1))
+        return tr["x"]["fn"]
 
-@pytest.mark.filterwarnings("ignore:can't resolve package")
-def test_independent():
-    from numpyro.contrib.tfp import distributions as tfd
+    res = f(0.0)
 
-    d = tfd.Independent(
-        tfd.Normal(jnp.zeros(10), jnp.ones(10)), reinterpreted_batch_ndims=1
-    )
-    assert d.event_shape == (10,)
-    assert d.batch_shape == ()
+    assert isinstance(res, TFPDistribution)
+    assert res.loc == 0
+    assert res.scale == 1
 
 
 @pytest.mark.filterwarnings("ignore:can't resolve package")
 def test_transformed_distributions():
-    from tensorflow_probability.substrates.jax import bijectors as tfb
-    from tensorflow_probability.substrates.jax.distributions import Normal as TFPNormal
-
-    from numpyro.contrib.tfp import distributions as tfd
+    from tensorflow_probability.substrates.jax import (
+        bijectors as tfb,
+        distributions as tfd,
+    )
 
     d = dist.TransformedDistribution(dist.Normal(0, 1), dist.transforms.ExpTransform())
-    d1 = tfd.TransformedDistribution(TFPNormal(0, 1), tfb.Exp())
-    d2 = dist.TransformedDistribution(
-        dist.Normal(0, 1), tfd.BijectorTransform(tfb.Exp())
-    )
+    d1 = tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())
     x = random.normal(random.PRNGKey(0), (1000,))
     d_x = d.log_prob(x).sum()
     d1_x = d1.log_prob(x).sum()
-    d2_x = d2.log_prob(x).sum()
     assert_allclose(d_x, d1_x)
-    assert_allclose(d_x, d2_x)
 
 
 @pytest.mark.filterwarnings("ignore:can't resolve package")
 def test_logistic_regression():
-    from numpyro.contrib.tfp import distributions as tfd
+    from tensorflow_probability.substrates.jax import distributions as tfd
 
     N, dim = 3000, 3
     num_warmup, num_samples = (1000, 1000)
     data = random.normal(random.PRNGKey(0), (N, dim))
     true_coefs = jnp.arange(1.0, dim + 1.0)
     logits = jnp.sum(true_coefs * data, axis=-1)
-    labels = tfd.Bernoulli(logits=logits)(rng_key=random.PRNGKey(1))
+    labels = tfd.Bernoulli(logits=logits).sample(seed=random.PRNGKey(1))
 
     def model(labels):
         coefs = numpyro.sample("coefs", tfd.Normal(jnp.zeros(dim), jnp.ones(dim)))
@@ -102,7 +79,7 @@ def model(labels):
 # TODO: remove after https://github.com/tensorflow/probability/issues/1072 is resolved
 @pytest.mark.filterwarnings("ignore:Explicitly requested dtype")
 def test_beta_bernoulli():
-    from numpyro.contrib.tfp import distributions as tfd
+    from tensorflow_probability.substrates.jax import distributions as tfd
 
     num_warmup, num_samples = (500, 2000)
 
@@ -114,7 +91,9 @@ def model(data):
         return p_latent
 
     true_probs = jnp.array([0.9, 0.1])
-    data = tfd.Bernoulli(true_probs)(rng_key=random.PRNGKey(1), sample_shape=(1000, 2))
+    data = tfd.Bernoulli(true_probs).sample(
+        seed=random.PRNGKey(1), sample_shape=(1000, 2)
+    )
     kernel = NUTS(model=model, trajectory_length=0.1)
     mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
     mcmc.run(random.PRNGKey(2), data)
@@ -240,14 +219,79 @@ def potential_fn(z):
 # test if sampling from tfp distributions works as expected using
 # numpyro sample function: numpyro.sample("name", dist) (bug)
 @pytest.mark.filterwarnings("ignore:can't resolve package")
+@pytest.mark.filterwarnings("ignore:Importing distributions")
 def test_sample_tfp_distributions():
-    from numpyro.contrib.tfp import distributions as tfd
+    from tensorflow_probability.substrates.jax import distributions as tfd
+
+    from numpyro.contrib.tfp.distributions import TFPDistribution
 
     # test no error raised
-    d = tfd.Normal(0, 1)
+    d = TFPDistribution[tfd.Normal](0, 1)
     with numpyro.handlers.seed(rng_seed=random.PRNGKey(0)):
         numpyro.sample("normal", d)
 
     # test intermediates are []
     value, intermediates = d(sample_intermediates=True, rng_key=random.PRNGKey(0))
     assert intermediates == []
+
+
+# test that sampling from unwrapped tensorflow_probability distributions works as
+# expected using numpyro.sample primitive
+@pytest.mark.parametrize(
+    "dist,args",
+    [
+        ["Bernoulli", (0,)],
+        ["Beta", (1, 1)],
+        ["Binomial", (10, 0)],
+        ["Categorical", ([0, 1, -1],)],
+        ["Cauchy", (0, 1)],
+        ["Dirichlet", ([1, 2, 0.5],)],
+        ["Exponential", (1,)],
+        ["InverseGamma", (1, 1)],
+        ["Normal", (0, 1)],
+        ["OrderedLogistic", ([0, 1], 0.5)],
+        ["Pareto", (1,)],
+    ],
+)
+def test_sample_unwrapped_tfp_distributions(dist, args):
+    from tensorflow_probability.substrates.jax import distributions as tfd
+
+    # test no error is raised
+    with numpyro.handlers.seed(rng_seed=random.PRNGKey(0)):
+        # since we import tfd inside the test, distributions have to be parametrized as
+        # strings, which is why we use getattr here
+        numpyro.sample("sample", getattr(tfd, dist)(*args))
+
+
+# test mixture distributions
+def test_sample_unwrapped_mixture_same_family():
+    from tensorflow_probability.substrates.jax import distributions as tfd
+
+    # test no error is raised
+    with numpyro.handlers.seed(rng_seed=random.PRNGKey(0)):
+        numpyro.sample(
+            "sample",
+            tfd.MixtureSameFamily(
+                mixture_distribution=tfd.Categorical(probs=[0.3, 0.7]),
+                components_distribution=tfd.Normal(
+                    loc=[-1.0, 1], scale=[0.1, 0.5]  # One for each component.
+                ),
+            ),
+        )
+
+
+# test that MCMC works with unwrapped tensorflow_probability distributions
+def test_mcmc_unwrapped_tfp_distributions():
+    from tensorflow_probability.substrates.jax import distributions as tfd
+
+    def model(y):
+        theta = numpyro.sample("p", tfd.Beta(1, 1))
+
+        with numpyro.plate("plate", y.size):
+            numpyro.sample("y", tfd.Bernoulli(probs=theta), obs=y)
+
+    mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=1000)
+    mcmc.run(random.PRNGKey(0), jnp.array([0, 0, 1, 1, 1]))
+    samples = mcmc.get_samples()
+
+    assert_allclose(jnp.mean(samples["p"]), 4 / 7, atol=0.05)