Repository: google/jax

Issue:
np.cumsum not implemented.
If the JAX version of np.cumsum were prioritized as a function to implement, I'd appreciate it.  I've implemented my own for now.

Thank you,
-David

Exception: Numpy function <function cumsum at 0x7f600be74a60> not yet implemented


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -370,6 +370,14 @@ def _reduce_window_sum(operand, window_dimensions, window_strides, padding):
       window_strides=tuple(window_strides), padding=padding,
       input_shape=operand.shape)
 
+def _reduce_window_prod(operand, window_dimensions, window_strides, padding):
+  init_value = _const(operand, 1)
+  jaxpr, consts = _reduction_jaxpr(mul, init_value)
+  return reduce_window_p.bind(
+      operand, init_value, jaxpr=jaxpr, consts=consts,
+      window_dimensions=tuple(window_dimensions),
+      window_strides=tuple(window_strides), padding=padding)
+
 def _reduce_window_max(operand, window_dimensions, window_strides, padding):
   return reduce_window_max_p.bind(
       operand, window_dimensions=tuple(window_dimensions),
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -882,6 +882,56 @@ def nan_reduction(a, axis=None, out=None, keepdims=False, **kwargs):
 nansum = _make_nan_reduction(onp.nansum, sum, 0, nan_if_all_nan=False)
 nanprod = _make_nan_reduction(onp.nanprod, prod, 1, nan_if_all_nan=False)
 
+
+def _make_cumulative_reduction(onp_reduction, window_reduce, init_val,
+                               squash_nan=False):
+  @_wraps(onp_reduction)
+  def cumulative_reduction(a, axis=None, dtype=None):
+    if axis is None or isscalar(a):
+      a = ravel(a)
+      axis = 0
+
+    a_shape = list(shape(a))
+    num_dims = len(a_shape)
+
+    if axis < 0:
+      axis = axis + num_dims
+    if axis < 0 or axis >= num_dims:
+      raise ValueError(
+          "axis {} is out of bounds for array of dimension {}".format(
+              axis, num_dims))
+
+    if squash_nan:
+      a = where(isnan(a), _constant_like(a, init_val), a)
+
+    if dtype:
+      a = lax.convert_element_type(a, dtype)
+
+    if a_shape[axis] == 0:
+      return a
+
+    padding = [(0, 0, 0)] * num_dims
+    padding[axis] = (a_shape[axis] - 1, 0, 0)
+    a = lax.pad(a, _constant_like(a, init_val), padding)
+    strides = [1] * num_dims
+    window_dims = [1] * num_dims
+    window_dims[axis] = a_shape[axis]
+    return window_reduce(
+       a, window_dims, strides, xla_bridge.get_xla_client().PaddingType.VALID)
+
+  return cumulative_reduction
+
+
+cumsum = _make_cumulative_reduction(
+  onp.cumsum, lax._reduce_window_sum, 0, squash_nan=False)
+cumprod = _make_cumulative_reduction(
+  onp.cumprod, lax._reduce_window_prod, 1, squash_nan=False)
+nancumsum = _make_cumulative_reduction(
+  onp.nancumsum, lax._reduce_window_sum, 0, squash_nan=True)
+nancumprod = _make_cumulative_reduction(
+  onp.nancumprod, lax._reduce_window_prod, 1, squash_nan=True)
+
+
 ### Array-creation functions
 
 @_wraps(onp.pad)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -529,6 +529,29 @@ def testRepeat(self, axis, shape, dtype, repeats, rng):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "op={}_shape=[{}]_axis={}_out_dtype={}".format(
+          op, jtu.format_shape_dtype_string(shape, dtype), axis, out_dtype),
+       "axis": axis, "shape": shape, "dtype": dtype, "out_dtype": out_dtype,
+       "rng": jtu.rand_default(), "lnp_op": getattr(lnp, op),
+       "onp_op": getattr(onp, op)}
+      for op in ["cumsum", "cumprod"]
+      # TODO(phawkins): replace both type lists with default_dtypes after a
+      # Jaxlib update includes
+      # https://github.com/google/jax/commit/86f5d189cf563b027c3cd00eea38072c003905c8
+      for dtype in [onp.float32, onp.int32]
+      for out_dtype in [onp.float32, onp.int32]
+      for shape in all_shapes
+      for axis in [None] + list(range(-len(shape), len(shape)))))
+  def testCumSumProd(self, axis, shape, dtype, out_dtype, onp_op, lnp_op, rng):
+    onp_fun = lambda arg: onp_op(arg, axis=axis, dtype=out_dtype)
+    lnp_fun = lambda arg: lnp_op(arg, axis=axis, dtype=out_dtype)
+
+    args_maker = lambda: [rng(shape, dtype)]
+
+    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}_m={}_n={}_k={}".format(
           onp.dtype(dtype).name, m, n, k),