Repository: google/jax

Issue:
Add autodiff support for `reduce_window`
I get the following exception trying to take the grad of a cumprod-using function, e.g.

```python
jax.jacrev(lambda x: jnp.cumprod(x, axis=-1))(np.random.randn(5))
```

```
File "jax/api.py", line 623, in batched_fun
    out_flat = batching.batch(jaxtree_fun, in_flat, in_axes_, out_axes)
  File "jax/interpreters/batching.py", line 45, in batch
    return batch_transform(fun, sz, in_dims, out_dim_dst).call_wrapped(in_vals)
  File "jax/linear_util.py", line 161, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "jax/api.py", line 501, in jacfun
    y, pullback = vjp(f_partial, *dyn_args)
  File "jax/api.py", line 1019, in vjp
    out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
  File "jax/interpreters/ad.py", line 105, in vjp
    out_primal, pval, jaxpr, consts = linearize(traceable, *primals)
  File "jax/interpreters/ad.py", line 94, in linearize
    jaxpr, out_pval, consts = pe.trace_to_jaxpr(jvpfun, in_pvals)
  File "jax/interpreters/partial_eval.py", line 400, in trace_to_jaxpr
    jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals)
  File "jax/linear_util.py", line 161, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  ... my code, ultimately calling np.cumprod ...

  File "jax/numpy/lax_numpy.py", line 1206, in cumulative_reduction
    return _cumulative_reduction(a, axis, dtype)
  File "jax/api.py", line 151, in f_jitted
    device_assignment=device_assignment)
  File "jax/core.py", line 675, in call_bind
    ans = full_lower(top_trace.process_call(primitive, f, tracers, params))
  File "jax/interpreters/ad.py", line 260, in process_call
    result = call_primitive.bind(f_jvp, pack(primals), nonzero_tangents, **params)
  File "jax/core.py", line 675, in call_bind
    ans = full_lower(top_trace.process_call(primitive, f, tracers, params))
  File "jax/interpreters/partial_eval.py", line 116, in process_call
    out_pv_const, consts = call_primitive.bind(fun, *in_consts, **params)
  File "jax/core.py", line 675, in call_bind
    ans = full_lower(top_trace.process_call(primitive, f, tracers, params))
  File "jax/interpreters/batching.py", line 135, in process_call
    val_out = call_primitive.bind(f, *vals, **params)
  File "jax/core.py", line 672, in call_bind
    ans = primitive.impl(f, *args, **params)
  File "jax/interpreters/xla.py", line 667, in _xla_call_impl
    *map(abstractify, args))
  File "jax/linear_util.py", line 213, in cached_fun
    ans, f_prev = cached_fun_body(f, args)
  File "jax/linear_util.py", line 210, in cached_fun_body
    return call(f, *args), f
  File "jax/interpreters/xla.py", line 679, in _xla_callable
    jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
  File "jax/linear_util.py", line 161, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "jax/numpy/lax_numpy.py", line 1201, in _cumulative_reduction
    a, window_dims, strides, xla_client.PaddingType.VALID)
  File "jax/lax/lax.py", line 939, in _reduce_window_prod
    window_strides=tuple(window_strides), padding=padding)
  File "jax/core.py", line 148, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "jax/interpreters/ad.py", line 251, in process_primitive
    .format(primitive))
NotImplementedError: Forward-mode differentiation rule for 'reduce_window' not implemented
```


Discussion:
It's easy enough to mimic TF's behavior here:
```
def cumprod_jvp(g, ans, x):
  return jnp.cumsum(g / x) * ans
cumprod = custom_transforms(jnp.cumprod)
defjvp(cumprod, cumprod_jvp)
```

This has two downsides:
a) it doesn't support the extra keyword arguments (dtype and axis). We need to extend `custom_transforms` a bit to allow that.
b) it doesn't work correctly if any entry in `x` is 0. Note that TF has the same bug (mishandling 0s). PyTorch does not have this bug because it falls back to a much more expensive quadratic algorithm if any entry is 0.


(There's also a clearly correct solution by rewriting `cumprod` using `lax.scan`, but it will most likely be slower than the current implementation.)
this is also something that we are interested in using in [pyhf](https://github.com/scikit-hep/pyhf) to compute hessians of likelihoods cc @kratsg @matthewfeickert

Resolving patch:
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
@@ -33,6 +33,7 @@
 import re
 import string
 import types
+from typing import Callable
 import warnings
 
 import numpy as onp
@@ -1569,7 +1570,59 @@ def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
   td = lax.div(nansum(a, axis, dtype=dtype, keepdims=keepdims), normalizer)
   return td
 
-def _make_cumulative_reduction(onp_reduction, window_reduce, init_val,
+
+# Parallel prefix-scan. See:
+# https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda
+# and
+# Blelloch, Guy E. 1990. "Prefix Sums and Their Applications.", Technical Report
+# CMU-CS-90-190, School of Computer Science, Carnegie Mellon University.
+#
+# Unlike the Blelloch algorithm, we use an out-of-place algorithm that uses 2n
+# space. This is somewhat wasteful if we are interested only in the output of
+# the forward pass, but more memory-efficient if we intend to differentiate
+# through the implementation of the scan.
+def _prescan_power_of_two(x, axis: int, op: Callable, unit):
+  n = x.shape[axis]
+  assert n != 0 and n & (n - 1) == 0, "n must be a power of 2"
+
+  # Upsweep
+  xs = []
+  for d in range(0, n.bit_length() - 1):
+    x1 = lax.slice_in_dim(x, 0, None, stride=2, axis=axis)
+    xs.append(x1)
+    x2 = lax.slice_in_dim(x, 1, None, stride=2, axis=axis)
+    x = op(x1, x2)
+  total = x
+
+  # Downsweep
+  x = full_like(total, unit)
+  pad_left = [(0, 0, 0)] * len(x.shape)
+  pad_left[axis] = (1, 0, 1)
+  pad_right = [(0, 0, 0)] * len(x.shape)
+  pad_right[axis] = (0, 1, 1)
+  for w in reversed(xs):
+    x1 = lax.pad(x, x.dtype.type(0), pad_right)
+    x2 = lax.pad(x, x.dtype.type(0), pad_left)
+    w = lax.pad(w, x.dtype.type(0), pad_left)
+    x = x1 + op(x2, w)
+
+  return x, total
+
+def _parallel_prefix_scan(x, axis: int, op: Callable, unit):
+  n = x.shape[axis]
+
+  # Pads to the next largest power of two
+  nbits = n.bit_length()
+  if n == (1 << (nbits - 1)):
+    nbits -= 1
+  padding = [(0, 0, 0)] * len(x.shape)
+  padding[axis] = (0, (1 << nbits) - n, 0)
+  x = lax.pad(x, x.dtype.type(unit), padding)
+  x, product = _prescan_power_of_two(x, axis, op, unit)
+  return concatenate((lax.slice_in_dim(x, 1, n, axis=axis), product), axis=axis)
+
+
+def _make_cumulative_reduction(onp_reduction, op, unit,
                                squash_nan=False):
   # We want to allow XLA to fuse the pad and reduce-window operators to
   # avoid materializing the padded output.
@@ -1592,7 +1645,7 @@ def _cumulative_reduction(a, axis, dtype):
               axis, num_dims))
 
     if squash_nan:
-      a = where(isnan(a), _constant_like(a, init_val), a)
+      a = where(isnan(a), _constant_like(a, unit), a)
 
     if not dtype and _dtype(a) == bool_:
       dtype = int_
@@ -1601,15 +1654,7 @@ def _cumulative_reduction(a, axis, 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_client.PaddingType.VALID)
+    return _parallel_prefix_scan(a, axis, op, unit)
 
   @_wraps(onp_reduction)
   def cumulative_reduction(a, axis=None, dtype=None):
@@ -1619,14 +1664,14 @@ def cumulative_reduction(a, axis=None, dtype=None):
 
 
 cumsum = _make_cumulative_reduction(
-  onp.cumsum, lax._reduce_window_sum, 0, squash_nan=False)
+  onp.cumsum, add, 0, squash_nan=False)
 cumprod = _make_cumulative_reduction(
-  onp.cumprod, lax._reduce_window_prod, 1, squash_nan=False)
+  onp.cumprod, multiply, 1, squash_nan=False)
 cumproduct = cumprod
 nancumsum = _make_cumulative_reduction(
-  onp.nancumsum, lax._reduce_window_sum, 0, squash_nan=True)
+  onp.nancumsum, add, 0, squash_nan=True)
 nancumprod = _make_cumulative_reduction(
-  onp.nancumprod, lax._reduce_window_prod, 1, squash_nan=True)
+  onp.nancumprod, multiply, 1, squash_nan=True)
 
 
 ### Array-creation functions


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
@@ -1104,7 +1104,8 @@ def attempt_sideeffect(x):
 
   @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),
+          op, jtu.format_shape_dtype_string(shape, dtype), axis,
+          out_dtype.__name__),
        "axis": axis, "shape": shape, "dtype": dtype, "out_dtype": out_dtype,
        "rng_factory": jtu.rand_default, "jnp_op": getattr(jnp, op),
        "onp_op": getattr(onp, op)}
@@ -1124,6 +1125,9 @@ def testCumSumProd(self, axis, shape, dtype, out_dtype, onp_op, jnp_op, rng_fact
     self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True,
                             tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
+    grad_dtypes = [onp.float32, onp.float64, onp.complex64, onp.complex128]
+    if dtype in grad_dtypes and out_dtype in grad_dtypes:
+      check_grads(jnp_fun, args_maker(), order=2, rtol=1e-2)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}_m={}_n={}_k={}".format(