Repository: google/jax

Issue:
Could we have support for scipy.stats.logistic?
In #324 we got support for `logit` and `expit` from `scipy.special`. This provides all the building blocks to also support `scipy.stats.logistic` (in particular, it’s members `cdf`, `pdf`, `ppf`, `logcdf` and `logpdf`). Having this would enable doing maximum-likelihood estimation of logistic regression using `scipy.optimize`.


Discussion:

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 pareto
 from . import t
 from . import uniform
+from . import logistic
diff --git a/jax/scipy/stats/logistic.py b/jax/scipy/stats/logistic.py
new file mode 100644
--- /dev/null
+++ b/jax/scipy/stats/logistic.py
@@ -0,0 +1,48 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import scipy.stats as osp_stats
+from jax.scipy.special import expit, logit
+
+from ... import lax
+from ...numpy.lax_numpy import _promote_args_inexact, _wraps
+
+
+@_wraps(osp_stats.logistic.logpdf, update_doc=False)
+def logpdf(x):
+  return lax.neg(x) - 2. * lax.log1p(lax.exp(lax.neg(x)))
+
+@_wraps(osp_stats.logistic.pdf, update_doc=False)
+def pdf(x):
+  return lax.exp(logpdf(x))
+
+@_wraps(osp_stats.logistic.ppf, update_doc=False)
+def ppf(x):
+  return logit(x)
+
+@_wraps(osp_stats.logistic.sf, update_doc=False)
+def sf(x):
+  return expit(lax.neg(x))
+
+@_wraps(osp_stats.logistic.isf, update_doc=False)
+def isf(x):
+  return -logit(x)
+
+@_wraps(osp_stats.logistic.cdf, update_doc=False)
+def cdf(x):
+  return expit(x)


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
@@ -211,6 +211,58 @@ def args_maker():
                             tol=1e-6)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
+  @genNamedParametersNArgs(1, jtu.rand_default)
+  def testLogisticCdf(self, rng_factory, shapes, dtypes):
+    rng = rng_factory()
+    scipy_fun = osp_stats.logistic.cdf
+    lax_fun = lsp_stats.logistic.cdf
+
+    def args_maker():
+      return list(map(rng, shapes, dtypes))
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                            tol=1e-6)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
+  @genNamedParametersNArgs(1, jtu.rand_default)
+  def testLogisticLogpdf(self, rng_factory, shapes, dtypes):
+    rng = rng_factory()
+    scipy_fun = osp_stats.logistic.logpdf
+    lax_fun = lsp_stats.logistic.logpdf
+
+    def args_maker():
+      return list(map(rng, shapes, dtypes))
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                            tol=1e-6)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
+  @genNamedParametersNArgs(1, jtu.rand_default)
+  def testLogisticPpf(self, rng_factory, shapes, dtypes):
+    rng = rng_factory()
+    scipy_fun = osp_stats.logistic.ppf
+    lax_fun = lsp_stats.logistic.ppf
+
+    def args_maker():
+      return list(map(rng, shapes, dtypes))
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                            tol=1e-6)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
+  @genNamedParametersNArgs(1, jtu.rand_default)
+  def testLogisticSf(self, rng_factory, shapes, dtypes):
+    rng = rng_factory()
+    scipy_fun = osp_stats.logistic.sf
+    lax_fun = lsp_stats.logistic.sf
+
+    def args_maker():
+      return list(map(rng, shapes, dtypes))
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
+                            tol=1e-6)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
   # TODO: currently it ignores the argument "shapes" and only tests dim=4
   @genNamedParametersNArgs(3, jtu.rand_default)
   def testMultivariateNormalLogPdf(self, rng_factory, shapes, dtypes):