Repository: google/jax
Issue: `vmap` fails for a reduction (e.g., `np.argmax`) over `bool`s

On Google Cloud Colab running the latest release of JAX I can find (0.1.22, last modified on 2019-07-20):

```python
!pip install --upgrade -q https://storage.googleapis.com/jax-wheels/cuda$(echo $CUDA_VERSION | sed -e 's/\\.//' -e 's/\\..*//')/jaxlib-0.1.22-cp36-none-linux_x86_64.whl
!pip install --upgrade -q jax

from __future__ import division
from __future__ import print_function

import jax
import jax.numpy as np
import numpy as onp

jax.config.update('jax_platform_name', 'cpu')

m = onp.random.random((5, 5))
m
# array([[0.05789924, 0.78492625, 0.14877936, 0.86753036, 0.71233691],
#        [0.58521529, 0.6402027 , 0.61487041, 0.7442337 , 0.57715207],
#        [0.29544111, 0.91596755, 0.17300919, 0.10794559, 0.46030352],
#        [0.06554419, 0.13804503, 0.72298234, 0.79693837, 0.04571761],
#        [0.8495216 , 0.70473064, 0.68557481, 0.57669368, 0.516764  ]])

np.argmax(m > 0.5, -1)  # DeviceArray([1, 0, 1, 2, 0], dtype=int32)
jax.vmap(np.argmax)(1.0 * (m > 0.5))  # DeviceArray([1, 0, 1, 2, 0], dtype=int32)
jax.vmap(np.argmax)(m > 0.5)
```
gives
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in get_primitive_batcher(p)
    231   try:
--> 232     return primitive_batchers[p]
    233   except KeyError:

KeyError: reduce

During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
10 frames
<ipython-input-32-f78d0b39613e> in <module>()
----> 1 jax.vmap(np.argmax)(m > 0.5)

/usr/local/lib/python3.6/dist-packages/jax/api.py in batched_fun(*args, **kwargs)
    491     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
    492     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
--> 493     out_flat = batching.batch(jaxtree_fun, in_flat, in_axes_, out_axes)
    494     return build_tree(out_tree(), out_flat)
    495 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dst)
     43   elif len(sizes) == 1:
     44     sz = sizes.pop()
---> 45     return batch_transform(fun, sz, in_dims, out_dim_dst).call_wrapped(in_vals)
     46   else:
     47     raise TypeError("got inconsistent map dimension sizes: {}".format(sizes))

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    145 
    146     del gen
--> 147     ans = self.f(*args, **dict(self.params, **kwargs))
    148     del args
    149     while stack:

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in argmax(a, axis)
   2011     a = ravel(a)
   2012     axis = 0
-> 2013   return _argminmax(max, a, axis)
   2014 
   2015 

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in _argminmax(op, a, axis)
   2028   idxs = onp.arange(a.shape[axis]).reshape(shape)
   2029   maxval = onp.iinfo(xla_bridge.canonicalize_dtype(idxs.dtype)).max
-> 2030   mask_idxs = where(lax._eq_meet(a, op(a, axis, keepdims=True)), idxs, maxval)
   2031   return min(mask_idxs, axis)
   2032 

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in reduction(a, axis, dtype, out, keepdims)
    956     if _dtype(a) != result_dtype:
    957       a = lax.convert_element_type(a, result_dtype)
--> 958     result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
    959     if keepdims:
    960       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)
    791     jaxpr, consts = _reduction_jaxpr(computation, init_value)
    792     return reduce_p.bind(operand, init_value, computation=computation,
--> 793                          jaxpr=jaxpr, consts=consts, dimensions=tuple(dimensions))
    794 
    795 def _reduction_jaxpr(computation, init_value):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    145 
    146     tracers = map(top_trace.full_raise, args)
--> 147     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    148     return full_lower(out_tracer)
    149 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    121     else:
    122       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
--> 123       batched_primitive = get_primitive_batcher(primitive)
    124       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    125       return BatchTracer(self, val_out, dim_out)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py in get_primitive_batcher(p)
    233   except KeyError:
    234     raise NotImplementedError(
--> 235         "Batching rule for '{}' not implemented".format(p))
    236 
    237 def defvectorized(prim):

NotImplementedError: Batching rule for 'reduce' not implemented
```

Changed files:
jax/numpy/lax_numpy.py, tests/lax_test.py, jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 14cc7842a..f8809ea4a 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -3255,7 +3255,7 @@ def _reduction_computation(c, jaxpr, consts, init_value):
 
 reduce_p = standard_primitive(_reduce_shape_rule, _input_dtype, 'reduce',
                               _reduce_translation_rule)
-# batching.primitive_batchers[reduce_p] = _reduce_batch_rule  # TODO(mattjj): test
+batching.primitive_batchers[reduce_p] = _reduce_batch_rule
 
 
 def _reduce_sum_shape_rule(operand, axes, input_shape):
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index a49f74c0f..904116ebb 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -991,6 +991,8 @@ def _reduction_dims(a, axis):
 
 def _reduction_init_val(a, init_val):
   a_dtype = xla_bridge.canonicalize_dtype(_dtype(a))
+  if a_dtype == 'bool':
+    return onp.array(init_val > 0, dtype=a_dtype)
   try:
     return onp.array(init_val, dtype=a_dtype)
   except OverflowError:
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 3a7f72de9..d55c54621 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1040,6 +1040,7 @@ class LaxTest(jtu.JaxTestCase):
       for init_val, op, dtypes in [
           (0, lax.add, default_dtypes),
           (1, lax.mul, default_dtypes),
+          (0, lax.max, all_dtypes), # non-monoidal
           (-onp.inf, lax.max, float_dtypes),
           (onp.iinfo(onp.int32).min, lax.max, [onp.int32]),
           # (onp.iinfo(onp.int64).min, lax.max, [onp.int64]),  # TODO fails
@@ -2591,14 +2592,15 @@ class LaxVmapTest(jtu.JaxTestCase):
     self._CheckBatching(op, 5, bdims, (shape,), dtype, rng)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_op={}_inshape={}_reducedims={}_bdims={}"
+      {"testcase_name": "_op={}_inshape={}_reducedims={}_initval={}_bdims={}"
        .format(op.__name__, jtu.format_shape_dtype_string(shape, dtype), dims,
-               bdims),
+               init_val, bdims),
        "op": op, "init_val": init_val, "shape": shape, "dtype": dtype,
        "dims": dims, "bdims": bdims, "rng": rng}
       for init_val, op, dtypes in [
           (0, lax.add, default_dtypes),
           (1, lax.mul, default_dtypes),
+          (0, lax.max, all_dtypes), # non-monoidal
           (-onp.inf, lax.max, float_dtypes),
           (onp.iinfo(onp.int32).min, lax.max, [onp.int32]),
           (onp.iinfo(onp.int64).min, lax.max, [onp.int64]),