Repository: google/jax

Issue:
np.cumsum(np.array([True])) raises a runtime error
```
import jax

a = jax.np.array([True])
b = jax.np.cumsum(a)
```
```
---------------------------------------------------------------------------
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:

11 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: Expected element type in shape to be arithmetic type for operation add; got PRED.: 

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
<ipython-input-3-94787de425b8> in <module>()
      2 
      3 a = jax.np.array([True])
----> 4 b = jax.np.cumsum(a)

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in cumulative_reduction(a, axis, dtype)
   1365   def cumulative_reduction(a, axis=None, dtype=None):
   1366     # jit doesn't support kwargs as static_args.
-> 1367     return _cumulative_reduction(a, axis, dtype)
   1368 
   1369   return cumulative_reduction

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
    148     _check_args(args_flat)
    149     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 150     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend)
    151     return tree_unflatten(out_tree(), out)
    152 

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, f, *args, **params)
    592   if top_trace is None:
    593     with new_sublevel():
--> 594       outs = primitive.impl(f, *args, **params)
    595   else:
    596     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    359   device = params['device']
    360   backend = params.get('backend', None)
--> 361   compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
    362   try:
    363     return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    207       fun.populate_stores(stores)
    208     else:
--> 209       ans = call(fun, *args)
    210       cache[key] = (ans, fun.stores)
    211     return ans

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, *abstract_args)
    390   xla_consts = _map(c.Constant, consts)
    391   xla_args = _xla_callable_args(c, abstract_args, tuple_args)
--> 392   out_nodes = jaxpr_subcomp(c, jaxpr, backend, axis_env, xla_consts, (), *xla_args)
    393   built = c.Build(c.Tuple(*out_nodes))
    394 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, freevars, *args)
    260       ans = rule(c, *in_nodes, **eqn.params)
    261     elif eqn.primitive in translations:
--> 262       ans = translations[eqn.primitive](c, *in_nodes, **eqn.params)
    263     elif eqn.primitive in reduction_translations:
    264       new_params = check_backend_params(eqn.params, backend)

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _reduce_window_sum_translation_rule(c, operand, window_dimensions, window_strides, padding, input_shape)
   3533   scalar = ShapedArray((), dtype)
   3534   return c.ReduceWindow(operand, c.Constant(onp.array(0, dtype)),
-> 3535                         xla.primitive_computation(add_p, scalar, scalar),
   3536                         window_dimensions, window_strides, padding)
   3537 

/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: Expected element type in shape to be arithmetic type for operation add; got PRED.: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```


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
@@ -3627,6 +3627,9 @@ def reduce_window(x, window_dimensions, window_strides, padding):
 
 def _reduce_window_sum_shape_rule(operand, window_dimensions, window_strides,
                                   padding, input_shape):
+  if not dtypes.issubdtype(operand.dtype, onp.number):
+    msg = "operand to reduce_window_sum must have a number dtype, got {}"
+    raise TypeError(msg.format(onp.dtype(operand.dtype).name))
   return _common_reduce_window_shape_rule(operand, window_dimensions,
                                          window_strides, padding)
 
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
@@ -1534,6 +1534,8 @@ def _cumulative_reduction(a, axis, dtype):
     if squash_nan:
       a = where(isnan(a), _constant_like(a, init_val), a)
 
+    if not dtype and _dtype(a) == bool_:
+      dtype = int_
     if dtype:
       a = lax.convert_element_type(a, dtype)
 
@@ -1553,7 +1555,6 @@ def _cumulative_reduction(a, axis, dtype):
   def cumulative_reduction(a, axis=None, dtype=None):
     # jit doesn't support kwargs as static_args.
     return _cumulative_reduction(a, axis, dtype)
-
   return cumulative_reduction
 
 


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
@@ -1075,7 +1075,7 @@ def test_single(m, args_maker, repeats, axis):
        "rng_factory": jtu.rand_default, "lnp_op": getattr(lnp, op),
        "onp_op": getattr(onp, op)}
       for op in ["cumsum", "cumprod"]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for out_dtype in default_dtypes
       for shape in all_shapes
       for axis in [None] + list(range(-len(shape), len(shape)))))