Repository: google/jax

Issue:
Support for scipy.stats.norm.logcdf
Hi there,

Adding support for `scipy.stats.norm.logcdf` would be very useful to me (for probit in the likelihood).

A workaround would be good too. I saw there is `erf`, but I could really do with a log version for numerical reasons.

Thanks!


Discussion:

Resolving patch:
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -20,6 +20,7 @@
 import scipy.special as osp_special
 
 from .. import lax
+from ..numpy import lax_numpy as np
 from ..numpy.lax_numpy import _wraps, asarray, _reduction_dims, _constant_like
 
 
@@ -55,4 +56,385 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   amax_singletons = dimadd(amax)
   out = lax.add(lax.log(lax.reduce(lax.exp(lax.sub(a, amax_singletons)),
                                    _constant_like(a, 0), lax.add, dims)), amax)
-  return dimadd(out) if keepdims else out
\ No newline at end of file
+  return dimadd(out) if keepdims else out
+
+
+# Normal distributions
+
+# Functions "ndtr" and "ndtri" are derived from calculations made in:
+# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+# In the following email exchange, the author gives his consent to redistribute
+# derived works under an Apache 2.0 license.
+#
+# From: Stephen Moshier <steve@moshier.net>
+# Date: Sat, Jun 9, 2018 at 2:36 PM
+# Subject: Re: Licensing cephes under Apache (BSD-like) license.
+# To: rif <rif@google.com>
+#
+#
+#
+# Hello Rif,
+#
+# Yes, Google may distribute Cephes files under the Apache 2 license.
+#
+# If clarification is needed, I do not favor BSD over other free licenses.
+# I would agree that Apache 2 seems to cover the concern you mentioned
+# about sublicensees.
+#
+# Best wishes for good luck with your projects!
+# Steve Moshier
+#
+#
+#
+# On Thu, 31 May 2018, rif wrote:
+#
+# > Hello Steve.
+# > My name is Rif. I work on machine learning software at Google.
+# >
+# > Your cephes software continues to be incredibly useful and widely used. I
+# > was wondering whether it would be permissible for us to use the Cephes code
+# > under the Apache 2.0 license, which is extremely similar in permissions to
+# > the BSD license (Wikipedia comparisons). This would be quite helpful to us
+# > in terms of avoiding multiple licenses on software.
+# >
+# > I'm sorry to bother you with this (I can imagine you're sick of hearing
+# > about this by now), but I want to be absolutely clear we're on the level and
+# > not misusing your important software. In former conversation with Eugene
+# > Brevdo (ebrevdo@google.com), you wrote "If your licensing is similar to BSD,
+# > the formal way that has been handled is simply to add a statement to the
+# > effect that you are incorporating the Cephes software by permission of the
+# > author." I wanted to confirm that (a) we could use the Apache license, (b)
+# > that we don't need to (and probably you don't want to) keep getting
+# > contacted about individual uses, because your intent is generally to allow
+# > this software to be reused under "BSD-like" license, and (c) you're OK
+# > letting incorporators decide whether a license is sufficiently BSD-like?
+# >
+# > Best,
+# >
+# > rif
+# >
+# >
+# >
+
+# log_ndtr uses different functions over the ranges
+# (-infty, lower](lower, upper](upper, infty)
+# Lower bound values were chosen by examining where the support of ndtr
+# appears to be zero, relative to scipy's (which is always 64bit). They were
+# then made more conservative just to be safe. (Conservative means use the
+# expansion more than we probably need to.)
+_LOGNDTR_FLOAT64_LOWER = onp.array(-20, onp.float64)
+_LOGNDTR_FLOAT32_LOWER = onp.array(-10, onp.float32)
+
+# Upper bound values were chosen by examining for which values of 'x'
+# Log[cdf(x)] is 0, after which point we need to use the approximation
+# Log[cdf(x)] = Log[1 - cdf(-x)] approx -cdf(-x). We chose a value slightly
+# conservative, meaning we use the approximation earlier than needed.
+_LOGNDTR_FLOAT64_UPPER = onp.array(8, onp.float64)
+_LOGNDTR_FLOAT32_UPPER = onp.array(5, onp.float32)
+
+
+def ndtr(x):
+  """Normal distribution function.
+
+  Returns the area under the Gaussian probability density function, integrated
+  from minus infinity to x:
+
+  ```
+                    1       / x
+     ndtr(x)  = ----------  |    exp(-0.5 t**2) dt
+                sqrt(2 pi)  /-inf
+
+              = 0.5 (1 + erf(x / sqrt(2)))
+              = 0.5 erfc(x / sqrt(2))
+  ```
+
+  Args:
+    x: An array of type `float32`, `float64`.
+
+  Returns:
+    ndtr: An array with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x` is not floating-type.
+  """
+  x = np.asarray(x)
+  dtype = lax._dtype(x)
+  if dtype not in (np.float32, np.float64):
+    raise TypeError(
+        "x.dtype={} is not supported, see docstring for supported types."
+        .format(dtype))
+  return _ndtr(x)
+
+
+def _ndtr(x):
+  """Implements ndtr core logic."""
+  dtype = lax._dtype(x).type
+  half_sqrt_2 = dtype(0.5) * onp.sqrt(2., dtype=dtype)
+  w = x * half_sqrt_2
+  z = lax.abs(w)
+  y = lax.select(lax.lt(z, half_sqrt_2),
+                      dtype(1.) + lax.erf(w),
+                      lax.select(lax.gt(w, dtype(0.)),
+                                      dtype(2.) - lax.erfc(z),
+                                      lax.erfc(z)))
+  return 0.5 * y
+
+
+def ndtri(p):
+  """The inverse of the CDF of the Normal distribution function.
+
+  Returns x such that the area under the pdf from minus infinity to x is equal
+  to p.
+
+  A piece-wise rational approximation is done for the function.
+  This is a port of the implementation in netlib.
+
+  Args:
+    p: an array of type `float32`, `float64`.
+
+  Returns:
+    x: an array with `dtype=p.dtype`.
+
+  Raises:
+    TypeError: if `p` is not floating-type.
+  """
+  x = np.asarray(p)
+  dtype = lax._dtype(p)
+  if dtype not in (np.float32, np.float64):
+    raise TypeError(
+        "x.dtype={} is not supported, see docstring for supported types."
+        .format(dtype))
+  return _ndtri(p)
+
+
+def _ndtri(p):
+  """Implements ndtri core logic."""
+
+  # Constants used in piece-wise rational approximations. Taken from the cephes
+  # library:
+  # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+  p0 = list(reversed([-5.99633501014107895267E1,
+                      9.80010754185999661536E1,
+                      -5.66762857469070293439E1,
+                      1.39312609387279679503E1,
+                      -1.23916583867381258016E0]))
+  q0 = list(reversed([1.0,
+                      1.95448858338141759834E0,
+                      4.67627912898881538453E0,
+                      8.63602421390890590575E1,
+                      -2.25462687854119370527E2,
+                      2.00260212380060660359E2,
+                      -8.20372256168333339912E1,
+                      1.59056225126211695515E1,
+                      -1.18331621121330003142E0]))
+  p1 = list(reversed([4.05544892305962419923E0,
+                      3.15251094599893866154E1,
+                      5.71628192246421288162E1,
+                      4.40805073893200834700E1,
+                      1.46849561928858024014E1,
+                      2.18663306850790267539E0,
+                      -1.40256079171354495875E-1,
+                      -3.50424626827848203418E-2,
+                      -8.57456785154685413611E-4]))
+  q1 = list(reversed([1.0,
+                      1.57799883256466749731E1,
+                      4.53907635128879210584E1,
+                      4.13172038254672030440E1,
+                      1.50425385692907503408E1,
+                      2.50464946208309415979E0,
+                      -1.42182922854787788574E-1,
+                      -3.80806407691578277194E-2,
+                      -9.33259480895457427372E-4]))
+  p2 = list(reversed([3.23774891776946035970E0,
+                      6.91522889068984211695E0,
+                      3.93881025292474443415E0,
+                      1.33303460815807542389E0,
+                      2.01485389549179081538E-1,
+                      1.23716634817820021358E-2,
+                      3.01581553508235416007E-4,
+                      2.65806974686737550832E-6,
+                      6.23974539184983293730E-9]))
+  q2 = list(reversed([1.0,
+                      6.02427039364742014255E0,
+                      3.67983563856160859403E0,
+                      1.37702099489081330271E0,
+                      2.16236993594496635890E-1,
+                      1.34204006088543189037E-2,
+                      3.28014464682127739104E-4,
+                      2.89247864745380683936E-6,
+                      6.79019408009981274425E-9]))
+
+  dtype = lax._dtype(p).type
+  shape = np.shape(p)
+
+  def _create_polynomial(var, coeffs):
+    """Compute n_th order polynomial via Horner's method."""
+    coeffs = onp.array(coeffs, dtype)
+    if not coeffs.size:
+      return np.zeros_like(var)
+    return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var
+
+
+  maybe_complement_p = np.where(p > dtype(-onp.expm1(-2.)), dtype(1.) - p, p)
+  # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
+  # later on. The result from the computation when p == 0 is not used so any
+  # number that doesn't result in NaNs is fine.
+  sanitized_mcp = np.where(
+      maybe_complement_p <= dtype(0.),
+      np.full(shape, dtype(0.5)),
+      maybe_complement_p)
+
+  # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
+  w = sanitized_mcp - dtype(0.5)
+  ww = lax.square(w)
+  x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)
+                              / _create_polynomial(ww, q0))
+  x_for_big_p *= -dtype(onp.sqrt(2. * onp.pi))
+
+  # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
+  # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
+  # arrays based on whether p < exp(-32).
+  z = lax.sqrt(dtype(-2.) * lax.log(sanitized_mcp))
+  first_term = z - lax.log(z) / z
+  second_term_small_p = (
+      _create_polynomial(dtype(1.) / z, p2) /
+      _create_polynomial(dtype(1.) / z, q2) / z)
+  second_term_otherwise = (
+      _create_polynomial(dtype(1.) / z, p1) /
+      _create_polynomial(dtype(1.) / z, q1) / z)
+  x_for_small_p = first_term - second_term_small_p
+  x_otherwise = first_term - second_term_otherwise
+
+  x = np.where(sanitized_mcp > dtype(onp.exp(-2.)),
+                      x_for_big_p,
+                      np.where(z >= dtype(8.0), x_for_small_p, x_otherwise))
+
+  x = np.where(p > dtype(1. - onp.exp(-2.)), x, -x)
+  infinity = np.full(shape, dtype(onp.inf))
+  x_nan_replaced = np.where(
+      p <= dtype(0.0), -infinity, np.where(p >= dtype(1.0), infinity, x))
+  return x_nan_replaced
+
+
+def log_ndtr(x, series_order=3):
+  """Log Normal distribution function.
+
+  For details of the Normal distribution function see `ndtr`.
+
+  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
+  using an asymptotic series. Specifically:
+  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
+    `log(1-x) ~= -x, x << 1`.
+  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
+    and take a log.
+  - For `x <= lower_segment`, we use the series approximation of erf to compute
+    the log CDF directly.
+
+  The `lower_segment` is set based on the precision of the input:
+
+  ```
+  lower_segment = { -20,  x.dtype=float64
+                  { -10,  x.dtype=float32
+  upper_segment = {   8,  x.dtype=float64
+                  {   5,  x.dtype=float32
+  ```
+
+  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:
+
+  ```
+     ndtr(x) = scale * (1 + sum) + R_N
+     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
+     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
+     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
+  ```
+
+  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
+  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).
+
+
+  Args:
+    x: an array of type `float32`, `float64`.
+    series_order: Positive Python `integer`. Maximum depth to
+      evaluate the asymptotic expansion. This is the `N` above.
+
+  Returns:
+    log_ndtr: an array with `dtype=x.dtype`.
+
+  Raises:
+    TypeError: if `x.dtype` is not handled.
+    TypeError: if `series_order` is a not Python `integer.`
+    ValueError:  if `series_order` is not in `[0, 30]`.
+  """
+  if not isinstance(series_order, int):
+    raise TypeError("series_order must be a Python integer.")
+  if series_order < 0:
+    raise ValueError("series_order must be non-negative.")
+  if series_order > 30:
+    raise ValueError("series_order must be <= 30.")
+
+  x = np.asarray(x)
+  dtype = lax._dtype(x)
+
+  if dtype == np.float64:
+    lower_segment = _LOGNDTR_FLOAT64_LOWER
+    upper_segment = _LOGNDTR_FLOAT64_UPPER
+  elif dtype == np.float32:
+    lower_segment = _LOGNDTR_FLOAT32_LOWER
+    upper_segment = _LOGNDTR_FLOAT32_UPPER
+  else:
+    raise TypeError("x.dtype={} is not supported.".format(onp.dtype(dtype)))
+
+  # The basic idea here was ported from:
+  #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+  # We copy the main idea, with a few changes
+  # * For x >> 1, and X ~ Normal(0, 1),
+  #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
+  #     which extends the range of validity of this function.
+  # * We use one fixed series_order for all of 'x', rather than adaptive.
+  # * Our docstring properly reflects that this is an asymptotic series, not a
+  #   Taylor series. We also provided a correct bound on the remainder.
+  # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
+  #   x=0. This happens even though the branch is unchosen because when x=0
+  #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
+  #   regardless of whether dy is finite. Note that the minimum is a NOP if
+  #   the branch is chosen.
+  return np.where(
+      lax.gt(x, upper_segment),
+      -_ndtr(-x),  # log(1-x) ~= -x, x << 1
+      np.where(lax.gt(x, lower_segment),
+                      lax.log(_ndtr(lax.max(x, lower_segment))),
+                      _log_ndtr_lower(lax.min(x, lower_segment),
+                                      series_order)))
+
+
+def _log_ndtr_lower(x, series_order):
+  """Asymptotic expansion version of `Log[cdf(x)]`, appropriate for `x<<-1`."""
+  dtype = lax._dtype(x).type
+  x_2 = lax.square(x)
+  # Log of the term multiplying (1 + sum)
+  log_scale = -dtype(0.5) * x_2 - lax.log(-x) - dtype(0.5 * onp.log(2. * onp.pi))
+  return log_scale + lax.log(_log_ndtr_asymptotic_series(x, series_order))
+
+
+def _log_ndtr_asymptotic_series(x, series_order):
+  """Calculates the asymptotic series used in log_ndtr."""
+  dtype = lax._dtype(x).type
+  if series_order <= 0:
+    return onp.array(1, dtype)
+  x_2 = lax.square(x)
+  even_sum = np.zeros_like(x)
+  odd_sum = np.zeros_like(x)
+  x_2n = x_2  # Start with x^{2*1} = x^{2*n} with n = 1.
+  for n in range(1, series_order + 1):
+    y = onp.array(_double_factorial(2 * n - 1), dtype) / x_2n
+    if n % 2:
+      odd_sum += y
+    else:
+      even_sum += y
+    x_2n *= x_2
+  return dtype(1.) + even_sum - odd_sum
+
+
+def _double_factorial(n):
+  """The double factorial function for small Python integer `n`."""
+  return onp.prod(onp.arange(n, 1, -2))
\ No newline at end of file
diff --git a/jax/scipy/stats/norm.py b/jax/scipy/stats/norm.py
--- a/jax/scipy/stats/norm.py
+++ b/jax/scipy/stats/norm.py
@@ -21,7 +21,7 @@
 
 from ... import lax
 from ...numpy.lax_numpy import _promote_args_like, _constant_like, _wraps
-
+from .. import special
 
 @_wraps(osp_stats.norm.logpdf)
 def logpdf(x, loc=0, scale=1):
@@ -32,6 +32,19 @@ def logpdf(x, loc=0, scale=1):
   quadratic = lax.div(lax.pow(lax.sub(x, loc), two), scale_sqrd)
   return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)
 
+
 @_wraps(osp_stats.norm.pdf)
 def pdf(x, loc=0, scale=1):
   return lax.exp(logpdf(x, loc, scale))
+
+
+@_wraps(osp_stats.norm.cdf)
+def cdf(x, loc=0, scale=1):
+  x, loc, scale = _promote_args_like(osp_stats.norm.cdf, x, loc, scale)
+  return special.ndtr(lax.div(lax.sub(x, loc), scale))
+
+
+@_wraps(osp_stats.norm.logcdf)
+def logcdf(x, loc=0, scale=1):
+  x, loc, scale = _promote_args_like(osp_stats.norm.logcdf, x, loc, scale)
+  return special.log_ndtr(lax.div(lax.sub(x, loc), scale))
\ No newline at end of file


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -49,21 +49,27 @@
 
 
 OpRecord = collections.namedtuple("OpRecord", ["name", "nargs", "dtypes", "rng",
-                                               "diff_modes", "test_name"])
+                                               "test_autodiff", "test_name"])
 
 
-def op_record(name, nargs, dtypes, rng, diff_modes, test_name=None):
+def op_record(name, nargs, dtypes, rng, test_grad, test_name=None):
   test_name = test_name or name
-  return OpRecord(name, nargs, dtypes, rng, diff_modes, test_name)
+  return OpRecord(name, nargs, dtypes, rng, test_grad, test_name)
 
 JAX_SPECIAL_FUNCTION_RECORDS = [
-    op_record("gammaln", 1, float_dtypes, jtu.rand_positive(), ["rev"]),
-    op_record("digamma", 1, float_dtypes, jtu.rand_positive(), []),
-    op_record("erf", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
-    op_record("erfc", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
-    op_record("erfinv", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
-    op_record("logit", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
-    op_record("expit", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
+    # TODO: digamma has no JVP implemented.
+    op_record("digamma", 1, float_dtypes, jtu.rand_positive(), False),
+    op_record("erf", 1, float_dtypes, jtu.rand_small_positive(), True),
+    op_record("erfc", 1, float_dtypes, jtu.rand_small_positive(), True),
+    op_record("erfinv", 1, float_dtypes, jtu.rand_small_positive(), True),
+    op_record("expit", 1, float_dtypes, jtu.rand_small_positive(), True),
+    # TODO: gammaln has slightly high error.
+    op_record("gammaln", 1, float_dtypes, jtu.rand_positive(), False),
+    # TODO: NaNs in gradient for logit.
+    op_record("logit", 1, float_dtypes, jtu.rand_small_positive(), False),
+    op_record("log_ndtr", 1, float_dtypes, jtu.rand_default(), True),
+    op_record("ndtri", 1, float_dtypes, jtu.rand_uniform(0., 1.), True),
+    op_record("ndtr", 1, float_dtypes, jtu.rand_default(), True),
 ]
 
 CombosWithReplacement = itertools.combinations_with_replacement
@@ -100,15 +106,15 @@ def lax_fun(array_to_reduce):
       {"testcase_name": jtu.format_test_name_suffix(
           rec.test_name, shapes, dtypes),
        "rng": rec.rng, "shapes": shapes, "dtypes": dtypes,
-       "modes": rec.diff_modes,
+       "test_autodiff": rec.test_autodiff,
        "scipy_op": getattr(osp_special, rec.name),
        "lax_op": getattr(lsp_special, rec.name)}
       for rec in JAX_SPECIAL_FUNCTION_RECORDS
       for shapes in CombosWithReplacement(all_shapes, rec.nargs)
       for dtypes in CombosWithReplacement(rec.dtypes, rec.nargs)))
-  def testScipySpecialFun(self, scipy_op, lax_op, rng, shapes, dtypes, modes):
+  def testScipySpecialFun(self, scipy_op, lax_op, rng, shapes, dtypes,
+                          test_autodiff):
     # TODO(mattjj): unskip this test combination when real() on tpu is improved
-    # TODO(mattjj): test autodiff
     if (FLAGS.jax_test_dut and FLAGS.jax_test_dut.startswith("tpu")
         and not shapes[0]):
       return absltest.unittest.skip("real() on scalar not supported on tpu")
@@ -119,5 +125,9 @@ def testScipySpecialFun(self, scipy_op, lax_op, rng, shapes, dtypes, modes):
                         check_dtypes=False)
     self._CompileAndCheck(lax_op, args_maker, check_dtypes=True)
 
+    if test_autodiff:
+      jtu.check_grads(lax_op, args, order=1, atol=1e-3, rtol=3e-3)
+
+
 if __name__ == "__main__":
   absltest.main()
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
@@ -134,6 +134,37 @@ def args_maker():
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
+
+  @genNamedParametersNArgs(3, jtu.rand_default())
+  def testNormLogCdf(self, rng, shapes, dtypes):
+    scipy_fun = osp_stats.norm.logcdf
+    lax_fun = lsp_stats.norm.logcdf
+
+    def args_maker():
+      x, loc, scale = map(rng, shapes, dtypes)
+      # clipping to ensure that scale is not too low
+      scale = onp.clip(onp.abs(scale), a_min=0.1, a_max=None)
+      return [x, loc, scale]
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
+
+  @genNamedParametersNArgs(3, jtu.rand_default())
+  def testNormCdf(self, rng, shapes, dtypes):
+    scipy_fun = osp_stats.norm.cdf
+    lax_fun = lsp_stats.norm.cdf
+
+    def args_maker():
+      x, loc, scale = map(rng, shapes, dtypes)
+      # clipping to ensure that scale is not too low
+      scale = onp.clip(onp.abs(scale), a_min=0.1, a_max=None)
+      return [x, loc, scale]
+
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
+
+
   @genNamedParametersNArgs(3, jtu.rand_default())
   def testUniformLogPdf(self, rng, shapes, dtypes):
     scipy_fun = osp_stats.uniform.logpdf