Repository: google/jax

Issue:
Cannot take gradient of np.prod
With the latest git version of jax, I get an error trying to `grad` a basic `np.prod` because some underlying `reduce` is not implemented. Perhaps this is a known limitation, but I could not find it in the documentation and it's what currently blocks me. Here is a minimal example:
```
import jax.numpy as np
from jax import grad

g = grad(lambda x: np.prod(x))
g(np.ones((10,8)))
```

> Traceback (most recent call last):
>   File "/home/guy/jax/jax/interpreters/ad.py", line 210, in process_primitive
>     jvp = primitive_jvps[primitive]
> KeyError: reduce
> ...
> NotImplementedError: Forward-mode differentiation rule for 'reduce' not implemented
> 


Discussion:
This should be easy enough to add. Note that the gradient of `np.sum` is implemented, so one possible workaround in the meantime is to compute `exp(sum(log(x)))`.

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -707,6 +707,8 @@ def _get_monoid_reducer(monoid_op, x):
   if (type(aval) is ConcreteArray) and aval.shape == ():
     if monoid_op is add:
       return aval.val == 0 and _reduce_sum
+    if monoid_op is mul:
+      return aval.val == 1 and _reduce_prod
     elif monoid_op is max:
       return aval.val == _get_max_identity(aval.dtype) and _reduce_max
     elif monoid_op is min:
@@ -735,6 +737,9 @@ def _get_min_identity(dtype):
 def _reduce_sum(operand, axes):
   return reduce_sum_p.bind(operand, axes=tuple(axes), input_shape=operand.shape)
 
+def _reduce_prod(operand, axes):
+  return reduce_prod_p.bind(operand, axes=tuple(axes))
+
 def _reduce_max(operand, axes):
   return reduce_max_p.bind(operand, axes=tuple(axes))
 
@@ -3029,6 +3034,54 @@ def _reduce_sum_transpose_rule(cotangent, input_shape, axes):
 batching.defreducer(reduce_sum_p)
 
 
+
+
+def _reduce_prod_shape_rule(operand, axes):
+  return tuple(onp.delete(operand.shape, axes))
+
+def _reduce_prod_translation_rule(c, operand, axes):
+  dtype = c.GetShape(operand).numpy_dtype()
+  scalar = xla_bridge.Shape.array_shape(dtype, ())
+  return c.Reduce(operand, c.Constant(onp.array(1, dtype)),
+                  xla.primitive_computation(mul_p, scalar, scalar),
+                  axes)
+
+def _reduce_prod_jvp_rule(tangent, operand, axes):
+  input_shape = onp.array(operand.shape)
+
+  n = onp.prod(input_shape[list(axes)])
+  non_axes = onp.delete(onp.arange(len(input_shape)), axes)
+
+  # Move the reduced axes to the front, and flatten them to 1D.
+  permutation = axes + tuple(non_axes)
+  new_shape = (n,) + tuple(input_shape[non_axes])
+  operand = reshape(operand, new_shape, permutation)
+  tangent = reshape(tangent, new_shape, permutation)
+
+  one = _const(operand, 1)
+  window_dims = [n] + [1] * len(non_axes)
+  window_strides = [1] * (len(non_axes) + 1)
+
+  # Form the partial products of all elements to the left and right of each
+  # element.
+  left_padding = [(n, -1, 0)] + [(0, 0, 0)] * len(non_axes)
+  right_padding = [(-1, n, 0)] + [(0, 0, 0)] * len(non_axes)
+  left_products = _reduce_window_prod(pad(operand, one, left_padding),
+                                      window_dims, window_strides,
+                                      xla_client.PaddingType.VALID)
+  right_products = _reduce_window_prod(pad(operand, one, right_padding),
+                                       window_dims, window_strides,
+                                       xla_client.PaddingType.VALID)
+
+  # Multiply partial products with the tangents and sum.
+  return _reduce_sum(mul(tangent, mul(left_products, right_products)), (0,))
+
+reduce_prod_p = standard_primitive(_reduce_prod_shape_rule, _input_dtype,
+                                   'reduce_prod', _reduce_prod_translation_rule)
+ad.defjvp(reduce_prod_p, _reduce_prod_jvp_rule)
+batching.defreducer(reduce_prod_p)
+
+
 def _reduce_chooser_shape_rule(operand, axes):
   return tuple(onp.delete(operand.shape, axes))
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1881,11 +1881,13 @@ def testTransposeGrad(self, shape, dtype, perm, rng):
        "dims": dims, "rng": rng}
       for init_val, op, dtypes in [
           (0, lax.add, inexact_dtypes),
+          (1, lax.mul, inexact_dtypes),
           (-onp.inf, lax.max, inexact_dtypes),
           (onp.inf, lax.min, inexact_dtypes),
       ]
       for dtype in dtypes
       for shape, dims in [
+          [(3, 4, 5), ()],
           [(3, 4, 5), (0,)],
           [(3, 4, 5), (1, 2)],
           [(3, 4, 5), (0, 2)],