Repository: pyro-ppl/numpyro

Issue:
Wrap distributions in scipy.stats for use with JAX
To make the internals of the `sampler` and `logpdf` methods visible to the tracer and compatible with jax, we need to:
 - Write the operations in terms of `jax.numpy` operations. 
 - We need to use jax's count based random number generator (via `PRNGKey`), and not the default used by scipy distributions which is numpy's globally mutable `mtrand`. Currently, the `PRNGKey` is being passed in through the `random_state` kwarg to the distribution's `_rvs` method.

Currently, only the normal distribution is wrapped. We can similarly wrap the following distributions (we may need to wrap over the samplers / rewrite the logic using jax operations):

**Continuous**
 - [x] beta @fehiepsi 
 - [x] cauchy @fehiepsi 
 - [x] dirichlet @fehiepsi 
 - [x] exponential @fehiepsi 
 - [x] uniform
 - [x] gamma  @fehiepsi 
 - [x] lognormal @fehiepsi 
 - [x] normal
 - [x] student-t  @fehiepsi 

**Discrete**
 - [x] bernoulli @neerajprad 
 - [x] categorical @neerajprad 
 - [x] binomial
 - [ ] poisson
 - [x] multinomial


Discussion:
@neerajprad Have you worked on this? I can't quickly add these distributions in case it is not overlapping with what you are doing. :)
> @neerajprad Have you worked on this? I can't quickly add these distributions in case it is not overlapping with what you are doing. :)

Feel free to work on this. I only took a brief look. I think the blocker is that we need to add a gamma sampler to jax first before we can support distributions like dirichlet and beta. What do you think?
@fehiepsi - Just add your name to the distribution that you are working on so that we don't end up doing double work! :) 
> we need to add a gamma sampler to jax first before we can support distributions like dirichlet and beta. 

Agree! I am mostly interested in seeing performance on hmm so I would like to support dirichlet distribution. Will try to implement a gamma sampler first (mainly to learn). :)
Hmm, looks like it is more complicated than I thought. Not because of understanding algorithms (Fritz already made it clear in his PR to pytorch), but because of my lacking knowledge on how jax/lax works (e.g. how to take gradient with shape parameters, how to implement if/else/while logics,...). I think that we'd need help from JAX devs on this.
I believe you'll need to define the sampler as a primitive operation and register its gradient separately, and not try to make the sampler itself differentiable. It is possible to do that but the documentation is a bit lacking - https://github.com/google/jax/issues/116. I think all we'll need is digamma to specify the gradient function which is already implemented in xla. What do you think? I can take a stab at it next week too.
Yeah, dwelling into `autograd` documentation, then searching for similar pattern in `jax.lax`, I come up with the following template, which works for some simple functions. ;)

```
from jax.interpreters import ad
from jax.lax import standard_unop, _float

def standard_gamma(x):
    # implement forward
    return x

def _standard_gamma_jvp_rule(g, ans, x):
    # implement backward
    return g * x

standard_gamma_p = standard_unop(_float, 'standard_gamma')
ad.defjvp2(standard_gamma_p, _standard_gamma_jvp_rule)
```

**Edit:** the above is not true (it is just applicable for XLA-compatible primitives), it seems that we need to define jvp_rule (forward-mode), transpose_rule (reverse-mode), batch_rule (vmap),... for custom primitives. I have implemented the sampler but the rules are still tricky for me.
@fehiepsi - We can ask the jax folks if this is already on their roadmap. In the meantime, feel free to add a PR for the sampler. For HMC at least, we just need the logpdf methods which can be wrapped over easily, and even with just a non-reparametrized sampler we can experiment with other algorithms.
Yes, I'll make a PR soon to not block us from adding more distributions.
No hurry, please take your time.

Resolving patch:
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -5,14 +5,16 @@
     expon,
     gamma,
     halfcauchy,
+    halfnorm,
     lognorm,
     norm,
     pareto,
     t,
     trunccauchy,
+    truncnorm,
     uniform
 )
-from numpyro.distributions.discrete import bernoulli, binom
+from numpyro.distributions.discrete import bernoulli, binom, poisson
 from numpyro.distributions.multivariate import categorical, dirichlet, multinomial
 
 __all__ = [
@@ -26,11 +28,14 @@
     'expon',
     'gamma',
     'halfcauchy',
+    'halfnorm',
     'lognorm',
     'multinomial',
     'norm',
     'pareto',
+    'poisson',
     't',
     'trunccauchy',
+    'truncnorm',
     'uniform',
 ]
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -6,10 +6,9 @@
 # Copyright (c) 2003-2019 SciPy Developers.
 # All rights reserved.
 
-
 import jax.numpy as np
 import jax.random as random
-import jax.scipy.stats as lsp_stats
+from jax.numpy.lax_numpy import _promote_dtypes
 from jax.scipy.special import digamma, gammaln, log_ndtr, ndtr, ndtri
 
 from numpyro.distributions import constraints
@@ -146,6 +145,33 @@ def _entropy(self):
         return np.log(2 * np.pi)
 
 
+class halfnorm_gen(jax_continuous):
+    _support_mask = constraints.positive
+
+    def _rvs(self):
+        return np.abs(random.normal(self._random_state, shape=self._size))
+
+    def _pdf(self, x):
+        # halfnorm.pdf(x) = sqrt(2/pi) * exp(-x**2/2)
+        return np.sqrt(2.0 / np.pi) * np.exp(-x * x / 2.0)
+
+    def _logpdf(self, x):
+        return 0.5 * np.log(2.0 / np.pi) - x * x / 2.0
+
+    def _cdf(self, x):
+        return norm._cdf(x) * 2 - 1.0
+
+    def _ppf(self, q):
+        return norm._ppf((1 + q) / 2.0)
+
+    def _stats(self):
+        return (np.sqrt(2.0 / np.pi), 1 - 2.0 / np.pi,
+                np.sqrt(2) * (4 - np.pi) / (np.pi - 2) ** 1.5, 8 * (np.pi - 3) / (np.pi - 2) ** 2)
+
+    def _entropy(self):
+        return 0.5 * np.log(np.pi / 2.0) + 0.5
+
+
 class lognorm_gen(jax_continuous):
     arg_constraints = {'s': constraints.positive}
     _support_mask = constraints.positive
@@ -164,19 +190,19 @@ def _logpdf(self, x, s):
                         -np.inf)
 
     def _cdf(self, x, s):
-        return ndtr(np.log(x) / s)
+        return norm._cdf(np.log(x) / s)
 
     def _logcdf(self, x, s):
-        return log_ndtr(np.log(x) / s)
+        return norm._logcdf(np.log(x) / s)
 
     def _ppf(self, q, s):
-        return np.exp(s * ndtri(q))
+        return np.exp(s * norm._ppf(q))
 
     def _sf(self, x, s):
-        return ndtr(-np.log(x) / s)
+        return norm._sf(np.log(x) / s)
 
     def _logsf(self, x, s):
-        return log_ndtr(-np.log(x) / s)
+        return norm._logsf(-np.log(x) / s)
 
     def _stats(self, s):
         p = np.exp(s * s)
@@ -196,6 +222,19 @@ class norm_gen(jax_continuous):
     def _rvs(self):
         return random.normal(self._random_state, shape=self._size)
 
+    def _pdf(self, x):
+        # norm.pdf(x) = exp(-x**2/2)/sqrt(2*pi)
+        return np.exp(-x**2 / 2.0) / np.sqrt(2 * np.pi)
+
+    def _logpdf(self, x):
+        return -(x ** 2 + np.log(2 * np.pi)) / 2.0
+
+    def _cdf(self, x):
+        return ndtr(x)
+
+    def _logcdf(self, x):
+        return log_ndtr(x)
+
     def _sf(self, x):
         return ndtr(-x)
 
@@ -214,12 +253,6 @@ def _stats(self):
     def _entropy(self):
         return 0.5 * (np.log(2 * np.pi) + 1)
 
-    def cdf(x, loc=0, scale=1):
-        return lsp_stats.norm.cdf(x, loc, scale)
-
-    def logcdf(x, loc=0, scale=1):
-        return lsp_stats.norm.logcdf(x, loc, scale)
-
 
 class pareto_gen(jax_continuous):
     arg_constraints = {'b': constraints.positive}
@@ -264,9 +297,6 @@ def _entropy(self, c):
         return 1 + 1.0 / c - np.log(c)
 
 
-pareto = pareto_gen(a=1.0, name="pareto")
-
-
 class t_gen(jax_continuous):
     arg_constraints = {'df': constraints.positive}
     _support_mask = constraints.real
@@ -320,14 +350,61 @@ def _rvs(self, a, b):
                            minval=np.arctan(a), maxval=np.arctan(b))
         return np.tan(u)
 
+    def _pdf(self, x, a, b):
+        return np.reciprocal((1 + x * x) * (np.arctan(b) - np.arctan(a)))
+
     def _logpdf(self, x, a, b):
         # trunc_pdf(x) = pdf(x) / (cdf(b) - cdf(a))
         #              = 1 / (1 + x^2) / (arctan(b) - arctan(a))
         normalizer = np.log(np.arctan(b) - np.arctan(a))
         return -(np.log(1 + x * x) + normalizer)
 
+
+class truncnorm_gen(jax_continuous):
+    # TODO: override _argcheck with the constraint that a < b
+
+    def _support(self, *args, **kwargs):
+        (a, b), loc, scale = self._parse_args(*args, **kwargs)
+        # TODO: make constraints.less_than and support a == -np.inf
+        if b == np.inf:
+            return constraints.greater_than((a - loc) * scale)
+        else:
+            return constraints.interval((a - loc) * scale, (b - loc) * scale)
+
+    def _rvs(self, a, b):
+        # We use inverse transform method:
+        # z ~ ppf(U), where U ~ Uniform(0, 1).
+        u = random.uniform(self._random_state, shape=self._size)
+        return self._ppf(u, a, b)
+
     def _pdf(self, x, a, b):
-        return np.reciprocal((1 + x * x) * (np.arctan(b) - np.arctan(a)))
+        delta = norm._sf(a) - norm._sf(b)
+        return norm._pdf(x) / delta
+
+    def _logpdf(self, x, a, b):
+        x, a, b = _promote_dtypes(x, a, b)
+        # XXX: consider to use norm._cdf(b) - norm._cdf(a) when a, b < 0
+        delta = norm._sf(a) - norm._sf(b)
+        return norm._logpdf(x) - np.log(delta)
+
+    def _cdf(self, x, a, b):
+        delta = norm._sf(a) - norm._sf(b)
+        return (norm._cdf(x) - norm._cdf(a)) / delta
+
+    def _ppf(self, q, a, b):
+        q, a, b = _promote_dtypes(q, a, b)
+        # XXX: consider to use norm._ppf(q * norm._cdf(b) + norm._cdf(a) * (1.0 - q))
+        # when a, b < 0
+        ppf = norm._isf(q * norm._sf(b) + norm._sf(a) * (1.0 - q))
+        return ppf
+
+    def _stats(self, a, b):
+        nA, nB = norm._cdf(a), norm._cdf(b)
+        d = nB - nA
+        pA, pB = norm._pdf(a), norm._pdf(b)
+        mu = (pA - pB) / d   # correction sign
+        mu2 = 1 + (a * pA - b * pB) / d - mu * mu
+        return mu, mu2, None, None
 
 
 class uniform_gen(jax_continuous):
@@ -354,8 +431,11 @@ def _entropy(self):
 expon = expon_gen(a=0.0, name='expon')
 gamma = gamma_gen(a=0.0, name='gamma')
 halfcauchy = halfcauchy_gen(a=0.0, name='halfcauchy')
+halfnorm = halfnorm_gen(a=0.0, name='halfnorm')
 lognorm = lognorm_gen(a=0.0, name='lognorm')
 norm = norm_gen(name='norm')
+pareto = pareto_gen(a=1.0, name="pareto")
 t = t_gen(name='t')
 trunccauchy = trunccauchy_gen(name='trunccauchy')
+truncnorm = truncnorm_gen(name='truncnorm')
 uniform = uniform_gen(a=0.0, b=1.0, name='uniform')
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -130,5 +130,42 @@ def _entropy(self, n, p):
         return np.sum(entr(vals), axis=0)
 
 
+class poisson_gen(jax_discrete):
+    arg_constraints = {'mu': constraints.positive}
+    _support_mask = constraints.nonnegative_integer
+
+    def _rvs(self, mu):
+        random_state = onp.random.RandomState(self._random_state)
+        sample = random_state.poisson(mu, self._size)
+        return device_put(sample)
+
+    def _logpmf(self, x, mu):
+        x, mu = _promote_dtypes(x, mu)
+        Pk = xlogy(x, mu) - gammaln(x + 1) - mu
+        return Pk
+
+    def _pmf(self, x, mu):
+        # poisson.pmf(k) = exp(-mu) * mu**k / k!
+        return np.exp(self._logpmf(x, mu))
+
+    def _cdf(self, x, mu):
+        raise NotImplementedError('Missing jax.scipy.special.pdtr')
+
+    def _sf(self, x, mu):
+        raise NotImplementedError('Missing jax.scipy.special.pdtrc')
+
+    def _ppf(self, q, mu):
+        raise NotImplementedError('Missing jax.scipy.special.pdtrk')
+
+    def _stats(self, mu):
+        var = mu
+        tmp = np.asarray(mu)
+        mu_nonzero = tmp > 0
+        g1 = np.where(mu_nonzero, np.sqrt(1.0 / tmp), np.inf)
+        g2 = np.where(mu_nonzero, 1.0 / tmp, np.inf)
+        return mu, var, g1, g2
+
+
 bernoulli = bernoulli_gen(b=1, name='bernoulli')
 binom = binom_gen(name='binom')
+poisson = poisson_gen(name='poisson')


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -32,9 +32,11 @@ def idfn(param):
     dist.expon,
     dist.gamma,
     dist.halfcauchy,
+    dist.halfnorm,
     dist.lognorm,
     dist.pareto,
     dist.trunccauchy,
+    dist.truncnorm,
     dist.norm,
     dist.t,
     dist.uniform,
@@ -76,6 +78,8 @@ def test_continuous_shape(jax_dist, loc, scale, prepend_shape):
     (dist.gamma, (np.array([-2., 3]),), np.array([1., -2])),
     (dist.halfcauchy, (), -1),
     (dist.halfcauchy, (), np.array([1., -2])),
+    (dist.halfnorm, (), -1),
+    (dist.halfnorm, (), np.array([1., -2])),
     (dist.lognorm, (-1,), -1),
     (dist.lognorm, (np.array([-2., 3]),), np.array([1., -2])),
     (dist.norm, (), np.inf),
@@ -86,6 +90,8 @@ def test_continuous_shape(jax_dist, loc, scale, prepend_shape):
     (dist.t, (np.array([-2., 3]),), np.array([1., np.nan])),
     (dist.trunccauchy, (), -1),
     (dist.trunccauchy, (), np.array([1., -2])),
+    (dist.truncnorm, (), -1),
+    (dist.truncnorm, (), np.array([1., -2])),
     (dist.uniform, (), -1),
     (dist.uniform, (), np.array([0.5, -2])),
 ], ids=idfn)
@@ -154,6 +160,8 @@ def test_multivariate_validate_args(jax_dist, valid_args, invalid_args, invalid_
     (dist.bernoulli, (np.array([0.3, 0.5]),)),
     (dist.binom, (10, 0.4)),
     (dist.binom, (np.array([10]), np.array([0.4, 0.3]))),
+    (dist.poisson, (1.,)),
+    (dist.poisson, (np.array([1., 4., 10.]),)),
 ], ids=idfn)
 @pytest.mark.parametrize('prepend_shape', [
     None,
@@ -178,6 +186,7 @@ def test_discrete_shape(jax_dist, dist_args, prepend_shape):
     (dist.bernoulli, (0.8,), (np.nan,), 2),
     (dist.binom, (10, 0.8), (-10, 0.8), -10),
     (dist.binom, (10, 0.8), (10, 1.1), -1),
+    (dist.poisson, (4.,), (-1.,), -1),
 ], ids=idfn)
 def test_discrete_validate_args(jax_dist, valid_args, invalid_args, invalid_sample):
     with validation_enabled():
@@ -195,12 +204,14 @@ def test_discrete_validate_args(jax_dist, valid_args, invalid_args, invalid_samp
     dist.expon,
     dist.gamma,
     dist.halfcauchy,
+    dist.halfnorm,
     dist.lognorm,
     dist.norm,
     dist.pareto,
     dist.t,
     pytest.param(dist.trunccauchy, marks=pytest.mark.xfail(
         reason='jvp rule for np.arctan is not yet available')),
+    dist.truncnorm,
     dist.uniform,
 ], ids=idfn)
 @pytest.mark.parametrize('loc, scale', [
@@ -243,11 +254,13 @@ def fn(args):
     dist.expon,
     dist.gamma,
     dist.halfcauchy,
+    dist.halfnorm,
     dist.lognorm,
     dist.norm,
     dist.pareto,
     dist.t,
     dist.trunccauchy,
+    dist.truncnorm,
     dist.uniform,
 ], ids=idfn)
 @pytest.mark.parametrize('loc_scale', [
@@ -328,6 +341,8 @@ def test_multivariate_discrete_logpmf(jax_dist, dist_args, shape):
     (dist.binom, (10, 0.4)),
     (dist.binom, (np.array([10]), np.array([0.4, 0.3]))),
     (dist.binom, (np.array([2, 5]), np.array([[0.4], [0.5]]))),
+    (dist.poisson, (4.,)),
+    (dist.poisson, (np.array([1., 4., 10.]),)),
 ], ids=idfn)
 @pytest.mark.parametrize('shape', [
     None,