# Buggy code
def _reduce_prod_jvp_rule(primals, tangents, *, axes):
    (operand,) = primals
    (tangent,) = tangents
    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)

    def _reduce_prod_tree(x, axis=0):
        """Reduce by repeatedly splitting the array and multiplying."""
        while x.shape[axis] > 1:
            n = x.shape[axis]
            n1 = (n + 1) // 2
            n2 = n - n1
            x1 = slice_in_dim(x, 0, n1)
            x2 = slice_in_dim(x, n1, None)
            if n2 != n1:
                paddings = [(0, 0, 0)] * len(x.shape)
                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)

    return api.jvp(_reduce_prod_tree, (operand,), (tangent,))


# Traceback (RuntimeError)
---------------------------------------------------------------------------
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

# Fixed code
def _reduce_prod_jvp_rule(primals, tangents, *, axes):
    (operand,) = primals
    (tangent,) = tangents
    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)

    def _reduce_prod_tree(x, axis=0):
        """Reduce by repeatedly splitting the array and multiplying."""
        while x.shape[axis] > 1:
            n = x.shape[axis]
            n1 = (n + 1) // 2
            n2 = n - n1
            x1 = slice_in_dim(x, 0, n1)
            x2 = slice_in_dim(x, n1, None)
            if n2 != n1:
                paddings = [(0, 0, 0)] * len(x.shape)
                paddings[axis] = (0, 1, 0)
                x2 = pad(x2, _const(x, 1), paddings)
            x = x1 * x2
        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,))