Repository: google/jax

Issue:
Require for Geometric Distribution
It seems that `jax` does not have geometric distribution in its random module like `scipy.stats`. Will it be added? I really need it.


Discussion:
We would happily accept pull requests to add a geometric distribution, but we have no immediate plans to work on it.
It is easy to simulate geometric random variables given a uniform RNG:

```np.ceil(np.log(np.random.rand()) / np.log1p(-p))```

has a `Geometric(p)` distribution. (Source: http://www.nrbook.com/devroye/Devroye_files/chapter_ten.pdf)

@gao462 I don't have time to package this up into a PR but I encourage you to do so if you end up using it.



Resolving patch:
diff --git a/jax/scipy/stats/__init__.py b/jax/scipy/stats/__init__.py
--- a/jax/scipy/stats/__init__.py
+++ b/jax/scipy/stats/__init__.py
@@ -26,3 +26,4 @@
 from . import t
 from . import uniform
 from . import logistic
+from . import geom
diff --git a/jax/scipy/stats/geom.py b/jax/scipy/stats/geom.py
new file mode 100644
--- /dev/null
+++ b/jax/scipy/stats/geom.py
@@ -0,0 +1,33 @@
+# Copyright 2020 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.
+
+import scipy.stats as osp_stats
+
+from ... import lax
+from ...numpy import lax_numpy as jnp
+from ...numpy._util import _wraps
+from ..special import xlogy, xlog1py
+
+@_wraps(osp_stats.geom.logpmf, update_doc=False)
+def logpmf(k, p, loc=0):
+    k, p, loc = jnp._promote_args_inexact("geom.logpmf", k, p, loc)
+    zero = jnp._constant_like(k, 0)
+    one = jnp._constant_like(k, 1)
+    x = lax.sub(k, loc)
+    log_probs = xlog1py(lax.sub(x, one), -p) + lax.log(p)
+    return jnp.where(lax.le(x, zero), -jnp.inf, log_probs)
+
+@_wraps(osp_stats.geom.pmf, update_doc=False)
+def pmf(k, p, loc=0):
+  return jnp.exp(logpmf(k, p, loc))


Test patch:
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
@@ -102,6 +102,23 @@ def args_maker():
                             tol=1e-4)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
+  @genNamedParametersNArgs(3, jtu.rand_default)
+  def testGeomLogPmf(self, rng_factory, shapes, dtypes):
+    rng = rng_factory(self.rng())
+    scipy_fun = osp_stats.geom.logpmf
+    lax_fun = lsp_stats.geom.logpmf
+
+    def args_maker():
+      x, logit, loc = map(rng, shapes, dtypes)
+      x = onp.floor(x)
+      p = expit(logit)
+      loc = onp.floor(loc)
+      return [x, p, loc]
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                            tol=1e-4)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
   @genNamedParametersNArgs(5, jtu.rand_positive)
   def testBetaLogPdf(self, rng_factory, shapes, dtypes):
     rng = rng_factory(self.rng())