Repository: google/jax
Issue: Batching broken for non-monoidal reducers

repro:

    >>> 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)


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

diff --git a/jax/lax.py b/jax/lax.py
index 00d3f480f..abc53a95a 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -111,7 +111,12 @@ def convert_element_type(operand, new_dtype):
     return operand
 
 def bitcast_convert_type(operand, new_dtype):
-  return bitcast_convert_type_p.bind(operand, new_dtype=new_dtype)
+  new_dtype = xla_bridge.canonicalize_dtype(new_dtype)
+  old_dtype = _dtype(operand)
+  if old_dtype != new_dtype:
+    return bitcast_convert_type_p.bind(operand, new_dtype=new_dtype)
+  else:
+    return operand
 
 def clamp(min, operand, max):
   return clamp_p.bind(min, operand, max)
@@ -256,8 +261,8 @@ def reduce(operand, init_value, computation, dimensions):
     return monoid_reducer(operand, dimensions)
   else:
     jaxpr, consts = _reduction_jaxpr(computation, init_value)
-    return reduce_p.bind(operand, init_value, jaxpr=jaxpr, consts=consts,
-                         dimensions=tuple(dimensions))
+    return reduce_p.bind(operand, init_value, computation=computation,
+                         jaxpr=jaxpr, consts=consts, dimensions=tuple(dimensions))
 
 def _reduction_jaxpr(computation, init_value):
   pval = _abstractify(init_value)
@@ -273,18 +278,26 @@ def _get_monoid_reducer(monoid_op, x):
       return aval.val == _get_max_identity(aval.dtype) and _reduce_max
     elif monoid_op is min:
       return aval.val == _get_min_identity(aval.dtype) and _reduce_min
+    elif monoid_op is bitwise_or and aval.dtype == onp.bool_:
+      return aval.val == _get_max_identity(aval.dtype) and _reduce_or
+    elif monoid_op is bitwise_and and aval.dtype == onp.bool_:
+      return aval.val == _get_min_identity(aval.dtype) and _reduce_and
 
 def _get_max_identity(dtype):
   if onp.issubdtype(dtype, onp.floating):
     return onp.array(-onp.inf, dtype)
   elif onp.issubdtype(dtype, onp.integer):
     return onp.array(onp.iinfo(dtype).min, dtype)
+  elif onp.issubdtype(dtype, onp.bool_):
+    return onp.array(False, onp.bool_)
 
 def _get_min_identity(dtype):
   if onp.issubdtype(dtype, onp.floating):
     return onp.array(onp.inf, dtype)
   elif onp.issubdtype(dtype, onp.integer):
     return onp.array(onp.iinfo(dtype).max, dtype)
+  elif onp.issubdtype(dtype, onp.bool_):
+    return onp.array(True, onp.bool_)
 
 def _reduce_sum(operand, axes):
   return reduce_sum_p.bind(operand, axes=tuple(axes), input_shape=operand.shape)
@@ -295,6 +308,12 @@ def _reduce_max(operand, axes):
 def _reduce_min(operand, axes):
   return reduce_min_p.bind(operand, axes=tuple(axes))
 
+def _reduce_or(operand, axes):
+  return reduce_or_p.bind(operand, axes=tuple(axes))
+
+def _reduce_and(operand, axes):
+  return reduce_and_p.bind(operand, axes=tuple(axes))
+
 def reduce_window(operand, init_value, computation, window_dimensions,
                   window_strides, padding):
   monoid_reducer = _get_monoid_window_reducer(computation, init_value)
@@ -1435,7 +1454,7 @@ def pad_batch_rule(batched_args, batch_dims, padding_config):
     padding_config.insert(operand_bdim, (0, 0, 0))
     return pad(operand, padding_value, padding_config), operand_bdim
   else:
-    raise NotImplementedError
+    raise NotImplementedError  # loop and stack
 
 pad_p = standard_primitive(pad_shape_rule, _input_dtype, 'pad')
 ad.deflinear(pad_p, pad_transpose)
@@ -1815,20 +1834,31 @@ ad.primitive_jvps[index_untake_p] = index_untake_jvp
 ad.primitive_transposes[index_untake_p] = index_untake_transpose_rule
 
 
-def reduce_shape_rule(operand, init_value, jaxpr, consts, dimensions):
+def reduce_shape_rule(operand, init_value, computation, jaxpr, consts, dimensions):
   return tuple(onp.delete(operand.shape, dimensions))
 
-def reduce_translation_rule(c, operand, init_value, jaxpr, consts, dimensions):
+def reduce_translation_rule(c, operand, init_value, computation, jaxpr, consts, dimensions):
   xla_computation = _reduction_computation(c, jaxpr, consts, init_value)
   return c.Reduce(operand, init_value, xla_computation, dimensions)
 
+def reduce_batch_rule(batched_args, batch_dims, computation, jaxpr, consts, dimensions):
+  operand, init_value = batched_args
+  operand_bdim, init_value_bdim = batch_dims
+  if init_value_bdim is None:
+    assert operand_bdim is not None
+    new_dimensions = [d + bool(d >= operand_bdim) for d in dimensions]
+    new_operand_bdim = operand_bdim - onp.sum(onp.less(dimensions, operand_bdim))
+    return reduce(operand, init_value, computation, new_dimensions), new_operand_bdim
+  else:
+    raise NotImplementedError  # loop and stack
+
 def _reduction_computation(c, jaxpr, consts, init_value):
   shape = c.GetShape(init_value)
   return xla.jaxpr_computation(jaxpr, consts, (), shape, shape)
 
 reduce_p = standard_primitive(reduce_shape_rule, _input_dtype, 'reduce',
                               reduce_translation_rule)
-batching.defreducer(reduce_p)
+# batching.primitive_batchers[reduce_p] = reduce_batch_rule  # TODO(mattjj): test
 
 
 def reduce_sum_shape_rule(operand, axes, input_shape):
@@ -1890,6 +1920,31 @@ ad.defjvp2(reduce_min_p, reduce_chooser_jvp_rule)
 batching.defreducer(reduce_min_p)
 
 
+def reduce_logical_shape_rule(operand, axes):
+  if operand.dtype != onp.bool_:
+    msg = "logical reduction requires operand dtype bool, got {}."
+    raise TypeError(msg.format(operand.dtype))
+  return tuple(onp.delete(operand.shape, axes))
+
+def reduce_logical_translation_rule(prim, identity, c, operand, axes):
+  scalar = xla_bridge.Shape.array_shape(onp.bool_, ())
+  return c.Reduce(operand, c.Constant(identity(onp.bool_)),
+                  xla.primitive_computation(prim, scalar, scalar), axes)
+
+reduce_or_translation_rule = partial(reduce_logical_translation_rule,
+                                     or_p, _get_max_identity)
+reduce_or_p = standard_primitive(reduce_logical_shape_rule, _fixed_dtype(onp.bool_),
+                                 'reduce_or', reduce_or_translation_rule)
+batching.defreducer(reduce_or_p)
+
+
+reduce_and_translation_rule = partial(reduce_logical_translation_rule,
+                                      and_p, _get_min_identity)
+reduce_and_p = standard_primitive(reduce_logical_shape_rule, _fixed_dtype(onp.bool_),
+                                 'reduce_and', reduce_and_translation_rule)
+batching.defreducer(reduce_and_p)
+
+
 def reduce_window_shape_rule(operand, init_value, jaxpr, consts,
                              window_dimensions, window_strides, padding):
   if operand.dtype != init_value.dtype:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 64905c311..ef92ae235 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -25,8 +25,7 @@ from .. import core
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray
 from ..interpreters.xla import DeviceArray
 from .. import lax
-from ..util import memoize
-from ..util import get_module_functions
+from ..util import memoize, partial, get_module_functions
 from ..lib import xla_bridge
 
 # To provide the same module-level names as Numpy, we need to redefine builtins
@@ -591,15 +590,16 @@ around = round
 ### Reducers
 
 
-def _make_reduction(np_fun, op, init_val):
+def _make_reduction(np_fun, op, init_val, preproc=None):
   """Creates reduction function given a binary operation and monoid identity."""
 
-  @_wraps(op)
+  @_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:
@@ -614,7 +614,6 @@ def _make_reduction(np_fun, op, init_val):
 
   return reduction
 
-
 def _reduction_dims(a, axis):
   if axis is None:
     return onp.arange(ndim(a))
@@ -625,7 +624,6 @@ def _reduction_dims(a, axis):
   else:
     raise TypeError("Unexpected type of axis argument: {}".format(type(axis)))
 
-
 def _reduction_init_val(a, init_val):
   a_dtype = xla_bridge.canonicalize_dtype(_dtype(a))
   try:
@@ -635,13 +633,14 @@ def _reduction_init_val(a, init_val):
     sign, iinfo = onp.sign(init_val), onp.iinfo(a_dtype)
     return onp.array(iinfo.min if sign < 0 else iinfo.max, dtype=a_dtype)
 
+_cast_to_bool = partial(lax.convert_element_type, new_dtype=onp.bool_)
 
 sum = _make_reduction(onp.sum, lax.add, 0)
 prod = _make_reduction(onp.prod, lax.mul, 1)
 max = _make_reduction(onp.max, lax.max, -onp.inf)
 min = _make_reduction(onp.min, lax.min, onp.inf)
-all = alltrue = _make_reduction(onp.all, logical_and, True)
-any = sometrue = _make_reduction(onp.any, logical_or, False)
+all = alltrue = _make_reduction(onp.all, lax.bitwise_and, True, _cast_to_bool)
+any = sometrue = _make_reduction(onp.any, lax.bitwise_or, False, _cast_to_bool)
 
 
 @_wraps(onp.mean)
diff --git a/tests/batching_test.py b/tests/batching_test.py
index 31f8a6cdb..4c15cc087 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -284,6 +284,13 @@ class BatchingTest(jtu.JaxTestCase):
     jacrev(func)(xs)  # don't crash
     jacfwd(func)(xs)  # don't crash
 
+  def testAny(self):
+    # test modeling the code in https://github.com/google/jax/issues/108
+
+    ans = vmap(np.any)(np.array([[True, False], [False, False]]))
+    expected = np.array([True, False])
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 04800e712..5d72ca4b3 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -142,8 +142,8 @@ JAX_REDUCER_RECORDS = [
 ]
 
 JAX_REDUCER_NO_DTYPE_RECORDS = [
-    op_record("all", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
-    op_record("any", 1, bool_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("all", 1, default_dtypes + bool_dtypes, all_shapes, jtu.rand_some_zero(), []),
+    op_record("any", 1, default_dtypes + bool_dtypes, all_shapes, jtu.rand_some_zero(), []),
     op_record("max", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
     op_record("min", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
 ]