Repository: google/jax

Issue:
support for scipy.stats.dirichlet
Registering interest in having support for this distribution (specifically logpdf() and rvs()).


Discussion:

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -32,7 +32,7 @@
 from . import numpy as np
 from . import tree_util
 from .api import jit, vmap
-from .numpy.lax_numpy import _constant_like
+from .numpy.lax_numpy import _constant_like, asarray
 from jax.lib import xla_bridge
 from jax import core
 
@@ -455,6 +455,30 @@ def _cauchy(key, shape, dtype):
   return lax.tan(lax.mul(pi, lax.sub(u, _constant_like(u, 0.5))))
 
 
+def dirichlet(key, alpha, shape=(), dtype=onp.float32):
+  """Sample Cauchy random values with given shape and float dtype.
+
+  Args:
+    key: a PRNGKey used as the random key.
+    alpha: an array-like with `alpha.shape[:-1]` broadcastable to `shape` and
+      used as the concentration parameter of the random variables.
+    shape: optional, a tuple of nonnegative integers representing the batch
+      shape (defaults to `alpha.shape[:-1]`).
+    dtype: optional, a float dtype for the returned values (default float32).
+
+  Returns:
+    A random array with the specified shape and dtype.
+  """
+  return _dirichlet(key, alpha, shape, dtype)
+
+@partial(jit, static_argnums=(2, 3))
+def _dirichlet(key, alpha, shape, dtype):
+  alpha = asarray(alpha, dtype)
+  shape = shape or alpha.shape[:-1]
+  gamma_samples = gamma(key, alpha, shape + alpha.shape[-1:], dtype)
+  return gamma_samples / np.sum(gamma_samples, axis=-1, keepdims=True)
+
+
 def exponential(key, shape=(), dtype=onp.float32):
   """Sample Exponential random values with given shape and float dtype.
 
diff --git a/jax/scipy/stats/dirichlet.py b/jax/scipy/stats/dirichlet.py
new file mode 100644
--- /dev/null
+++ b/jax/scipy/stats/dirichlet.py
@@ -0,0 +1,45 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+import scipy.stats as osp_stats
+
+from ... import lax
+from ...numpy.lax_numpy import _promote_args_like, _constant_like, _wraps, all, sum
+from ..special import gammaln, xlogy
+
+
+def _is_simplex(x):
+    x_sum = sum(x, axis=-1)
+    return all(x > 0, axis=-1) & (x_sum <= 1) & (x_sum > 1 - 1e-6)
+
+
+@_wraps(osp_stats.dirichlet.logpdf)
+def logpdf(x, alpha):
+    args = (onp.ones((0,), lax.dtype(x)), onp.ones((1,), lax.dtype(alpha)))
+    to_dtype = lax.dtype(osp_stats.dirichlet.logpdf(*args))
+    x, alpha = [lax.convert_element_type(arg, to_dtype) for arg in (x, alpha)]
+    one = _constant_like(x, 1)
+    normalize_term = sum(gammaln(alpha), axis=-1) - gammaln(sum(alpha, axis=-1))
+    log_probs = lax.sub(sum(xlogy(lax.sub(alpha, one), x), axis=-1), normalize_term)
+    return where(_is_simplex(x), log_probs, -inf)
+
+
+@_wraps(osp_stats.dirichlet.pdf)
+def pdf(x, alpha):
+  return lax.exp(logpdf(x, alpha))


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -212,6 +212,26 @@ def testCauchy(self, dtype):
     for samples in [uncompiled_samples, compiled_samples]:
       self._CheckKolmogorovSmirnovCDF(samples, scipy.stats.cauchy().cdf)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_alpha={}_{}".format(alpha, dtype),
+       "alpha": alpha, "dtype": onp.dtype(dtype).name}
+      for alpha in [[0.2, 1., 5.]]
+      for dtype in [onp.float32, onp.float64]))
+  @jtu.skip_on_devices("tpu")  # TODO(phawkins): re-enable
+  def testDirichlet(self, alpha, dtype):
+    key = random.PRNGKey(0)
+    rand = lambda key, alpha: random.dirichlet(key, alpha, (10000,), dtype)
+    crand = api.jit(rand)
+
+    uncompiled_samples = rand(key, alpha)
+    compiled_samples = crand(key, alpha)
+
+    for samples in [uncompiled_samples, compiled_samples]:
+      self.assertAllClose(samples.sum(-1), onp.ones(10000, dtype=dtype), check_dtypes=True)
+      alpha_sum = sum(alpha)
+      for i, a in enumerate(alpha):
+        self._CheckKolmogorovSmirnovCDF(samples[..., i], scipy.stats.beta(a, alpha_sum - a).cdf)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}".format(dtype), "dtype": onp.dtype(dtype).name}
       for dtype in [onp.float32, onp.float64]))
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -93,6 +93,21 @@ def args_maker():
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
+  @genNamedParametersNArgs(2, jtu.rand_positive())
+  def testDirichletLogPdf(self, rng, shapes, dtypes):
+    scipy_fun = osp_stats.cauchy.logpdf
+    lax_fun = lsp_stats.cauchy.logpdf
+    dim = 4
+    shapes = (shapes[0] + (dim,), shapes[1] + (dim,))
+
+    def args_maker():
+      x, alpha = map(rng, shapes, dtypes)
+      x = x / onp.sum(x, axis=-1, keepdims=True)
+      return [x, alpha]
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
   @genNamedParametersNArgs(3, jtu.rand_positive())
   def testExponLogPdf(self, rng, shapes, dtypes):
     scipy_fun = osp_stats.expon.logpdf