# Buggy code
def _reduction_init_val(a, init_val):
    a_dtype = xla_bridge.canonicalize_dtype(_dtype(a))
    try:
        return onp.array(init_val, dtype=a_dtype)
    except OverflowError:
        assert onp.issubdtype(a_dtype, onp.integer)
        sign, iinfo = onp.sign(init_val), onp.iinfo(a_dtype)
        return onp.array(iinfo.min if sign < 0 else iinfo.max, dtype=a_dtype)


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

# Fixed code
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:
        assert onp.issubdtype(a_dtype, onp.integer)
        sign, iinfo = onp.sign(init_val), onp.iinfo(a_dtype)
        return onp.array(iinfo.min if sign < 0 else iinfo.max, dtype=a_dtype)