Repository: google/jax

Issue:
Bitwise XOR missing as a `lax.reduce` operation
The `lax.reduce` function includes bitwise AND and OR as possible reduction operations, but does not include XOR. Likewise, there is no specialized XOR-reduction function in JAX.


Discussion:
Thanks for the report – I think it would be reasonable to add a `reduce_xor` primitive, modeled after `reduce_or` and `reduce_and`. Are you interested in contributing?
No -- I don't know how to add that. Are there bitwise AND and OR reductions
already? Or are there only those reductions for logical operations?

-- Jeremiah Willcock


On Wed, Jun 8, 2022 at 10:57 AM Jake Vanderplas ***@***.***>
wrote:

> Thanks for the report – I think it would be reasonable to add a reduce_xor
> primitive, modeled after reduce_or and reduce_and. Are you interested in
> contributing?
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/google/jax/issues/11005#issuecomment-1150223086>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAX34KHHQ2MDWYQJWWKYVSLVODNHXANCNFSM5YBQ2K4Q>
> .
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>

Yes, there are bitwise AND and OR reductions currently. I may be confused, because I thought you had referred to the existing bitwise AND and OR reductions in your initial comment. Is there something I'm not understanding about your request?
I was looking at _get_monoid_reducer in
https://jax.readthedocs.io/en/latest/_modules/jax/_src/lax/lax.html#reduce;
is that relevant to what reduction ops are available?

-- Jeremiah Willcock


On Wed, Jun 8, 2022 at 11:18 AM Jake Vanderplas ***@***.***>
wrote:

> Yes, there are bitwise AND and OR reductions currently. I may be confused,
> because I thought you had referred to the existing bitwise AND and OR
> reductions in your initial comment. Is there something I'm not
> understanding about your request?
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/google/jax/issues/11005#issuecomment-1150242523>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAX34KH6PIZSBMAYWB5FO5TVODPZBANCNFSM5YBQ2K4Q>
> .
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>

Yes, that's relevant here. And you'll see the list of currently available specialized reducers here: https://github.com/google/jax/blob/879a94a05405a84d677d0ab049cfefe0169dc22b/jax/_src/lax/lax.py#L970-L990
Note that even without the specialization, you can currently compute an xor reduction manually:
```python
from jax import lax
import jax.numpy as jnp
result = lax.reduce(jnp.arange(6, dtype='uint8'), jnp.uint8(0), lax.bitwise_xor, (0,))
print(result)
# 1
```

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -983,6 +983,8 @@ def _get_monoid_reducer(monoid_op: Callable,
       return np.equal(aval.val, _get_max_identity(dtype)) and _reduce_or
     elif monoid_op is bitwise_and and dtype == np.bool_:
       return np.equal(aval.val, _get_min_identity(dtype)) and _reduce_and
+    elif monoid_op is bitwise_xor and dtype == np.bool_:
+      return np.equal(aval.val, _get_max_identity(dtype)) and _reduce_xor
     elif monoid_op is max:
       return np.equal(aval.val, _get_max_identity(dtype)) and _reduce_max
     elif monoid_op is min:
@@ -1023,6 +1025,8 @@ def _reduce_or(operand: Array, axes: Sequence[int]) -> Array:
 def _reduce_and(operand: Array, axes: Sequence[int]) -> Array:
   return reduce_and_p.bind(operand, axes=tuple(axes))
 
+def _reduce_xor(operand: Array, axes: Sequence[int]) -> Array:
+  return reduce_xor_p.bind(operand, axes=tuple(axes))
 
 def sort(operand: Union[Array, Sequence[Array]], dimension: int = -1,
          is_stable: bool = True, num_keys: int = 1) -> Union[Array, Tuple[Array, ...]]:
@@ -3618,6 +3622,12 @@ def _reduce_logical_shape_rule(operand, *, axes):
 batching.defreducer(reduce_and_p)
 
 
+reduce_xor_p = standard_primitive(
+    _reduce_logical_shape_rule, _fixed_dtype(np.bool_), 'reduce_xor',
+    weak_type_rule=_strip_weak_type)
+batching.defreducer(reduce_xor_p)
+
+
 def _unary_reduce_lower(reducer, unit_factory, ctx, x, *, axes):
   aval_out, = ctx.avals_out
   dtype = aval_out.dtype
@@ -3639,6 +3649,8 @@ def _unary_reduce_lower(reducer, unit_factory, ctx, x, *, axes):
                                          lambda dtype: np.array(False, dtype)))
 mlir.register_lowering(reduce_and_p, partial(_unary_reduce_lower, mhlo.AndOp,
                                           lambda dtype: np.array(True, dtype)))
+mlir.register_lowering(reduce_xor_p, partial(_unary_reduce_lower, mhlo.XorOp,
+                                          lambda dtype: np.array(False, dtype)))
 mlir.register_lowering(reduce_min_p, partial(_unary_reduce_lower, mlir.min_mhlo,
                                          _get_min_identity))
 mlir.register_lowering(reduce_max_p, partial(_unary_reduce_lower, mlir.max_mhlo,
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -981,6 +981,7 @@ def _unexpected_primitive(p: core.Primitive, *args, **kwargs):
     "igamma_grad_a",
     "random_gamma_grad",
     "reduce_precision",
+    "reduce_xor",
     "schur",
     "closed_call",
     "unreachable",
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -182,6 +182,7 @@
   reduce_precision_p as reduce_precision_p,
   reduce_prod_p as reduce_prod_p,
   reduce_sum_p as reduce_sum_p,
+  reduce_xor_p as reduce_xor_p,
   regularized_incomplete_beta_p as regularized_incomplete_beta_p,
   rem as rem,
   rem_p as rem_p,


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1748,6 +1748,44 @@ def testReduce(self, op, init_val, shape, dtype, dims):
     args_maker = lambda: [rng(shape, dtype)]
     self._CompileAndCheck(fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_op={}_shape={}_reducedims={}_initval={}_prim={}"
+       .format(op.__name__, shape, dims, init_val, prim),
+       "op": op, "init_val": init_val, "shape": shape, "dims": dims, "prim": prim}
+      for init_val, op , prim in [
+        (True, lax.bitwise_and, jax.lax.reduce_and_p),
+        (False, lax.bitwise_or, jax.lax.reduce_or_p),
+        (False, lax.bitwise_xor, jax.lax.reduce_xor_p),
+      ]
+      for shape, dims in [
+          [(3, 4, 5), (0,)], [(3, 4, 5), (1, 2)],
+          [(3, 4, 5), (0, 2)], [(3, 4, 5), (0, 1, 2)]
+      ]))
+  def testReduceBoolean(self, op, init_val, shape, dims, prim):
+    def reference_fun(operand, init_value):
+      np_op = getattr(np, op.__name__)
+      return np_op.reduce(operand, axis=dims, initial=init_val)
+
+    dtype = np.bool_
+    rng = jtu.rand_bool(self.rng())
+    init_val = np.asarray(init_val, dtype=dtype)
+    fun = lambda operand, init_val: lax.reduce(operand, init_val, op, dims)
+    args_maker = lambda: [rng(shape, dtype), init_val]
+    self._CompileAndCheck(fun, args_maker)
+    self._CheckAgainstNumpy(reference_fun, fun, args_maker)
+
+    # recheck with a static init_val
+    fun = lambda operand: lax.reduce(operand, init_val, op, dims)
+    reference_fun = partial(reference_fun, init_value=init_val)
+    args_maker = lambda: [rng(shape, dtype)]
+    self._CompileAndCheck(fun, args_maker)
+    self._CheckAgainstNumpy(reference_fun, fun, args_maker)
+
+    # check that the correct monoid reducer primitive is used inside the
+    # jaxpr. This requires the init_val (monoid identity element) to be static
+    jaxpr = jax.make_jaxpr(fun)(rng(shape, dtype))
+    self.assertEqual(jaxpr.eqns[0].primitive, prim)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}.{}_arr_weak_type={}_init_weak_type={}"
        .format(op_namespace.__name__, op, arr_weak_type, init_weak_type),