Repository: google/jax

Issue:
Runtime error when taking grad of prod of take with empty indices
When taking grad/JVP of `lambda params: np.prod(np.take(params, []))`, a runtime error is produced.

Jax version: 0.1.55

Repro:
```python
jax.grad(lambda params: np.prod(np.take(params, np.array([], np.int32), axis=0)))(np.ones(6))
```
Error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in primitive_computation(prim, *avals, **params)
    179   try:
--> 180     return c.Build()
    181   except RuntimeError as e:

23 frames
/usr/local/lib/python3.6/dist-packages/jax/lib/xla_bridge.py in Build(self, *args, **kwargs)
    256     return super(_JaxComputationBuilder, self).Build(
--> 257         *args, **kwargs)
    258 

/usr/local/lib/python3.6/dist-packages/jaxlib/xla_client.py in Build(self, root, backend)
    729     else:
--> 730       return Computation(self._builder.Build(), backend=backend)
    731 

RuntimeError: Invalid argument: Padding result in negative size for dimension 0: 

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
<ipython-input-2-26b18538a688> in <module>()
----> 1 jax.grad(lambda params: np.prod(np.take(params, np.array([], np.int32), axis=0)))(np.ones(6))

/usr/local/lib/python3.6/dist-packages/jax/api.py in grad_f(*args, **kwargs)
    345   @wraps(fun, docstr=docstr, argnums=argnums)
    346   def grad_f(*args, **kwargs):
--> 347     _, g = value_and_grad_f(*args, **kwargs)
    348     return g
    349 

/usr/local/lib/python3.6/dist-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    400     f_partial, dyn_args = _argnums_partial(f, argnums, args)
    401     if not has_aux:
--> 402       ans, vjp_py = vjp(f_partial, *dyn_args)
    403     else:
    404       ans, vjp_py, aux = vjp(f_partial, *dyn_args, has_aux=True)

/usr/local/lib/python3.6/dist-packages/jax/api.py in vjp(fun, *primals, **kwargs)
   1255   if not has_aux:
   1256     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1257     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1258     out_tree = out_tree()
   1259   else:

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    105 def vjp(traceable, primals, has_aux=False):
    106   if not has_aux:
--> 107     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    108   else:
    109     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     94   _, in_tree = tree_flatten(((primals, primals), {}))
     95   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
---> 96   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
     97   pval_primals, pval_tangents = tree_unflatten(out_tree(), out_pvals)
     98   aval_primals, const_primals = unzip2(pval_primals)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, **kwargs)
    341   with new_master(JaxprTrace) as master:
    342     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 343     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    344     assert not env
    345     del master

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    151 
    152     del gen
--> 153     ans = self.f(*args, **dict(self.params, **kwargs))
    154     del args
    155     while stack:

<ipython-input-2-26b18538a688> in <lambda>(params)
----> 1 jax.grad(lambda params: np.prod(np.take(params, np.array([], np.int32), axis=0)))(np.ones(6))

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in reduction(a, axis, dtype, out, keepdims)
   1127     if _dtype(a) != result_dtype:
   1128       a = lax.convert_element_type(a, result_dtype)
-> 1129     result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
   1130     if keepdims:
   1131       shape_with_singletons = lax.subvals(shape(a), zip(dims, (1,) * len(dims)))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in reduce(operand, init_value, computation, dimensions)
    874   monoid_reducer = _get_monoid_reducer(computation, init_value)
    875   if monoid_reducer:
--> 876     return monoid_reducer(operand, dimensions)
    877   else:
    878     jaxpr, consts = _reduction_jaxpr(computation, _abstractify(init_value))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _reduce_prod(operand, axes)
    925 
    926 def _reduce_prod(operand, axes):
--> 927   return reduce_prod_p.bind(operand, axes=tuple(axes))
    928 
    929 def _reduce_max(operand, axes):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    153 
    154     tracers = map(top_trace.full_raise, args)
--> 155     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    156     if self.multiple_results:
    157       return map(full_lower, out_tracer)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    220           "Forward-mode differentiation rule for '{}' not implemented"
    221           .format(primitive))
--> 222     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    223     if primitive.multiple_results:
    224       return [JVPTracer(self, x, t) for x, t in zip(primal_out, tangent_out)]

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in standard_jvp(jvprules, primitive, primals, tangents, **params)
    319 def standard_jvp(jvprules, primitive, primals, tangents, **params):
    320   val_out = primitive.bind(*primals, **params)
--> 321   tangents_out = [rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
    322                   if rule is not None and t is not zero]
    323   return val_out, reduce(add_tangents, tangents_out, zero)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in <listcomp>(.0)
    320   val_out = primitive.bind(*primals, **params)
    321   tangents_out = [rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
--> 322                   if rule is not None and t is not zero]
    323   return val_out, reduce(add_tangents, tangents_out, zero)
    324 

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _reduce_prod_jvp_rule(tangent, operand, axes)
   3410   left_padding = [(n, -1, 0)] + [(0, 0, 0)] * len(non_axes)
   3411   right_padding = [(-1, n, 0)] + [(0, 0, 0)] * len(non_axes)
-> 3412   left_products = _reduce_window_prod(pad(operand, one, left_padding),
   3413                                       window_dims, window_strides,
   3414                                       xla_client.PaddingType.VALID)

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in pad(operand, padding_value, padding_config)
    634   operator.
    635   """
--> 636   return pad_p.bind(operand, padding_value, padding_config=tuple(padding_config))
    637 
    638 def rev(operand, dimensions):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    150     top_trace = find_top_trace(args)
    151     if top_trace is None:
--> 152       return self.impl(*args, **kwargs)
    153 
    154     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    138   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    139   abstract_args = map(abstractify, args)
--> 140   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
    141   return compiled_fun(*args)
    142 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
    150   else:
    151     handle_result = aval_to_result_handler(aval_out)
--> 152   built_c = primitive_computation(prim, *abstract_args, **params)
    153   compiled = built_c.Compile(compile_options=xb.get_compile_options(),
    154                              backend=xb.get_backend(backend))

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in primitive_computation(prim, *avals, **params)
    183            "This is a bug in JAX's shape-checking rules; please report it!\n"
    184            "https://github.com/google/jax/issues\n")
--> 185     raise RuntimeError(msg)
    186 
    187 def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):

RuntimeError: Invalid argument: Padding result in negative size for dimension 0: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```

This doesn't happen when taking a JVP of just `np.take`, i.e.
```
jax.jvp(lambda params: np.take(params, np.array([], np.int32), axis=0), (np.ones(6),), (np.ones(6),))
```


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4284,9 +4284,9 @@ def _reduce_prod_tree(x, axis=0):
         paddings[axis] = (0, 1, 0)
         x2 = pad(x2, _const(x, 1), paddings)
       x = x1 * x2
-    shape = list(x.shape)
-    del shape[axis]
-    return reshape(x, shape)
+    if x.shape[axis] == 0:
+      return full(input_shape[non_axes], _one(x))
+    return squeeze(x, (axis,))
 
   return api.jvp(_reduce_prod_tree, (operand,), (tangent,))
 


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -651,6 +651,7 @@ def testTransposeGrad(self, shape, dtype, perm, rng_factory):
           [(3, 4, 5), (0, 2)],
           [(3, 4, 5), (0, 1, 2)],
           [(3, 1), (1,)],
+          [(3, 0, 5), (1,)],
       ]))
   def testReduceGrad(self, op, init_val, shape, dtype, dims, rng_factory):
     rng = rng_factory(self.rng())
@@ -664,7 +665,8 @@ def testReduceGrad(self, op, init_val, shape, dtype, dims, rng_factory):
     eps = (1.0 if dtypes.finfo(dtype).bits == 16 and op is lax.add else
            1e-1 if dtype == dtypes.bfloat16 else
            1e-2 if dtypes.finfo(dtype).bits == 32 else None)
-    check_grads(reduce, (operand,), 2, ["fwd", "rev"], tol, tol, eps)
+    if op not in (lax.max, lax.min) or all(d > 0 for d in shape):
+      check_grads(reduce, (operand,), 2, ["fwd", "rev"], tol, tol, eps)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}_dtype={}_padding={}"
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
@@ -469,7 +469,7 @@ def testNotImplemented(self):
   def testOp(self, np_op, jnp_op, rng_factory, shapes, dtypes, check_dtypes,
              tolerance, inexact):
     np_op = jtu.ignore_warning(category=RuntimeWarning,
-                                message="invalid value.*")(np_op)
+                               message="invalid value.*")(np_op)
 
     rng = rng_factory(self.rng())
     args_maker = self._GetArgsMaker(rng, shapes, dtypes, np_arrays=False)