Repository: pyro-ppl/numpyro

Issue:
Implement JVP rules for cumsum/cumprod
Currently, cumsum/cumprod does not have jvp rule in jax yet. And these operators are necessary for simplex constraint. We should make custom primitives to support these operators.


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -26,7 +26,7 @@
 from jax.scipy.special import expit, logit
 
 from numpyro.distributions import constraints
-from numpyro.distributions.util import sum_rightmost
+from numpyro.distributions.util import cumprod, cumsum, sum_rightmost
 
 
 def clipped_expit(x):
@@ -177,7 +177,7 @@ def __call__(self, x):
         x = x - np.log(x.shape[-1] - np.arange(x.shape[-1]))
         # convert to probabilities (relative to the remaining) of each fraction of the stick
         z = clipped_expit(x)
-        z1m_cumprod = np.cumprod(1 - z, axis=-1)
+        z1m_cumprod = cumprod(1 - z)
         pad_width = [(0, 0)] * x.ndim
         pad_width[-1] = (0, 1)
         z_padded = np.pad(z, pad_width, mode="constant", constant_values=1.)
@@ -187,7 +187,7 @@ def __call__(self, x):
 
     def inv(self, y):
         y_crop = y[..., :-1]
-        z1m_cumprod = 1 - y_crop.cumsum(-1)
+        z1m_cumprod = 1 - cumsum(y_crop)
         # hence x = logit(z) = log(z / (1 - z)) = y[::-1] / z1m_cumprod
         x = np.log(y_crop / z1m_cumprod)
         return x + np.log(x.shape[-1] - np.arange(x.shape[-1]))
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -4,7 +4,7 @@
 import scipy.special as osp_special
 
 import jax.numpy as np
-from jax import canonicalize_dtype, jit, lax, random, vmap
+from jax import canonicalize_dtype, custom_transforms, jit, lax, random, vmap
 from jax.core import Primitive
 from jax.interpreters import ad, partial_eval, xla
 from jax.numpy.lax_numpy import _promote_args_like, _promote_shapes
@@ -278,6 +278,24 @@ def binary_cross_entropy_with_logits(x, y):
     return np.clip(x, 0) + np.log1p(np.exp(-np.abs(x))) - x * y
 
 
+@custom_transforms
+def cumsum(x):
+    return np.cumsum(x, axis=-1)
+
+
+ad.defjvp(cumsum.primitive, lambda g, x: np.cumsum(g, axis=-1))
+
+
+@custom_transforms
+def cumprod(x):
+    return np.cumprod(x, axis=-1)
+
+
+# XXX this implementation does not address the case x=0, hence the result in that case will be nan
+# Ref: https://stackoverflow.com/questions/40916955/how-to-compute-gradient-of-cumprod-safely
+ad.defjvp2(cumprod.primitive, lambda g, ans, x: np.cumsum(g / x, axis=-1) * ans)
+
+
 def promote_shapes(*args, shape=()):
     # adapted from lax.lax_numpy
     if len(args) < 2 and not shape:


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -395,17 +395,8 @@ def test_biject_to(constraint, shape):
     assert np.shape(actual) == batch_shape
     if len(shape) == transform.event_dim:
         if constraint is constraints.simplex:
-            # compare against pytorch
-            if shape == (1,):
-                expected = -1.3968685
-            elif shape == (3,):
-                expected = -7.139747
-            elif shape == (5,):
-                expected = -11.871904
-            inv_expected = -expected
-            # TODO compare against numerical values when jvp rule for `cumsum`, `cumprod` is available
-            # expected = np.linalg.slogdet(jax.jacobian(transform)(x)[::-1, :])[1]
-            # inv_expected = np.linalg.slogdet(jax.jacobian(transform.inv)(y)[:, ::-1])[1]
+            expected = np.linalg.slogdet(jax.jacobian(transform)(x)[:-1, :])[1]
+            inv_expected = np.linalg.slogdet(jax.jacobian(transform.inv)(y)[:, :-1])[1]
         else:
             expected = np.log(np.abs(grad(transform)(x)))
             inv_expected = np.log(np.abs(grad(transform.inv)(y)))
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -8,7 +8,7 @@
 from jax.scipy.special import expit
 from jax.util import partial
 
-from numpyro.distributions.util import binary_cross_entropy_with_logits, standard_gamma, xlog1py, xlogy
+from numpyro.distributions.util import binary_cross_entropy_with_logits, cumprod, cumsum, standard_gamma, xlog1py, xlogy
 
 _zeros = partial(lax.full_like, fill_value=0)
 
@@ -81,6 +81,31 @@ def test_binary_cross_entropy_with_logits(x, y):
     assert_allclose(actual, expect, rtol=1e-6)
 
 
+@pytest.mark.parametrize('shape', [
+    (3,),
+    (5, 4),
+])
+def test_cumsum_jac(shape):
+    rng = random.PRNGKey(0)
+    x = random.normal(rng, shape=shape)
+    expected = grad(lambda x: np.sum(x * (np.arange(x.shape[-1]) + 1.)[::-1]))(x)
+    actual = grad(lambda x: np.sum(cumsum(x)))(x)
+    assert_allclose(expected, actual)
+
+
+@pytest.mark.parametrize('shape', [
+    (3,),
+    (5, 3),
+])
+def test_cumprod_jac(shape):
+    rng = random.PRNGKey(0)
+    x = random.uniform(rng, shape=shape)
+    expected = grad(lambda x: np.sum(x[..., 0] + x[..., 0] * x[..., 1]
+                                     + x[..., 0] * x[..., 1] * x[..., 2]))(x)
+    actual = grad(lambda x: np.sum(cumprod(x)))(x)
+    assert_allclose(expected, actual, rtol=2e-7)
+
+
 @pytest.mark.parametrize('alpha, shape', [
     (1., ()),
     (1., (2,)),