# Buggy code
def _make_reduction(np_fun, op, init_val):
    """Creates reduction function given a binary operation and monoid identity."""

    @_wraps(op)
    def reduction(a, axis=None, dtype=None, out=None, keepdims=False):
        if out is not None:
            raise ValueError("reduction does not support `out` argument.")

        a = a if isinstance(a, ndarray) else asarray(a)
        dims = _reduction_dims(a, axis)
        result_dtype = _dtype(np_fun(onp.ones((), dtype=dtype or _dtype(a))))
        if _dtype(a) != result_dtype:
            a = lax.convert_element_type(a, result_dtype)
        result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
        if keepdims:
            shape_with_singletons = lax.subvals(shape(a), zip(dims, (1,) * len(dims)))
            result = lax.reshape(result, shape_with_singletons)
        if dtype and onp.dtype(dtype) != onp.dtype(result_dtype):
            result = lax.convert_element_type(result, dtype)
        return result

    return reduction


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def _make_reduction(np_fun, op, init_val, preproc=None):
    """Creates reduction function given a binary operation and monoid identity."""

    @_wraps(np_fun)
    def reduction(a, axis=None, dtype=None, out=None, keepdims=False):
        if out is not None:
            raise ValueError("reduction does not support `out` argument.")

        a = a if isinstance(a, ndarray) else asarray(a)
        a = preproc(a) if preproc else a
        dims = _reduction_dims(a, axis)
        result_dtype = _dtype(np_fun(onp.ones((), dtype=dtype or _dtype(a))))
        if _dtype(a) != result_dtype:
            a = lax.convert_element_type(a, result_dtype)
        result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
        if keepdims:
            shape_with_singletons = lax.subvals(shape(a), zip(dims, (1,) * len(dims)))
            result = lax.reshape(result, shape_with_singletons)
        if dtype and onp.dtype(dtype) != onp.dtype(result_dtype):
            result = lax.convert_element_type(result, dtype)
        return result

    return reduction