Specialize ufunc.reduce for monoidal binary ufuncs.

Author: jakevdp

jax/_src/numpy/reductions.py

@@ -192,6 +192,11 @@ def _cast_to_bool(operand: ArrayLike) -> Array:
 def _cast_to_numeric(operand: ArrayLike) -> Array:
   return promote_dtypes_numeric(operand)[0]
 
+def _require_integer(operand: ArrayLike) -> Array:
+  arr = lax_internal.asarray(operand)
+  if not dtypes.isdtype(arr, ("bool", "integral")):
+    raise ValueError(f"integer argument required; got dtype={arr.dtype}")
+  return arr
 
 def _ensure_optional_axes(x: Axis) -> Axis:
   def force(x):
@@ -652,6 +657,63 @@ def any(a: ArrayLike, axis: Axis = None, out: None = None,
   return _reduce_any(a, axis=_ensure_optional_axes(axis), out=out,
                      keepdims=keepdims, where=where)
 
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_bitwise_and(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                        out: None = None, keepdims: bool = False,
+                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  arr = lax_internal.asarray(a)
+  init_val = np.array(-1, dtype=dtype or arr.dtype)
+  return _reduction(arr, "reduce_bitwise_and", None, lax.bitwise_and, init_val, preproc=_require_integer,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_bitwise_or(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                        out: None = None, keepdims: bool = False,
+                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  return _reduction(a, "reduce_bitwise_or", None, lax.bitwise_or, 0, preproc=_require_integer,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_bitwise_xor(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                        out: None = None, keepdims: bool = False,
+                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  return _reduction(a, "reduce_bitwise_xor", None, lax.bitwise_xor, 0, preproc=_require_integer,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_logical_and(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                        out: None = None, keepdims: bool = False,
+                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  return _reduction(a, "reduce_logical_and", None, lax.bitwise_and, True, preproc=_cast_to_bool,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_logical_or(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                       out: None = None, keepdims: bool = False,
+                       initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  return _reduction(a, "reduce_logical_or", None, lax.bitwise_or, False, preproc=_cast_to_bool,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
+@partial(api.jit, static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
+def _reduce_logical_xor(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
+                        out: None = None, keepdims: bool = False,
+                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
+  return _reduction(a, "reduce_logical_xor", None, lax.bitwise_xor, False, preproc=_cast_to_bool,
+                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
+                    initial=initial, where_=where)
+
+
 def amin(a: ArrayLike, axis: Axis = None, out: None = None,
         keepdims: bool = False, initial: ArrayLike | None = None,
         where: ArrayLike | None = None) -> Array:

jax/_src/numpy/ufuncs.py

@@ -31,7 +31,7 @@
 from jax._src.custom_derivatives import custom_jvp
 from jax._src.lax import lax
 from jax._src.lax import other as lax_other
-from jax._src.typing import Array, ArrayLike, DTypeLike
+from jax._src.typing import Array, ArrayLike
 from jax._src.numpy.util import (
    check_arraylike, promote_args, promote_args_inexact,
    promote_args_numeric, promote_dtypes_inexact, promote_dtypes_numeric,
@@ -1221,7 +1221,7 @@ def multiply(x: ArrayLike, y: ArrayLike, /) -> Array:
   x, y = promote_args("multiply", x, y)
   return lax.mul(x, y) if x.dtype != bool else lax.bitwise_and(x, y)
 
-@binary_ufunc(identity=-1)
+@binary_ufunc(identity=-1, reduce=reductions._reduce_bitwise_and)
 def bitwise_and(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the bitwise AND operation elementwise.
 
@@ -1250,7 +1250,7 @@ def bitwise_and(x: ArrayLike, y: ArrayLike, /) -> Array:
   """
   return lax.bitwise_and(*promote_args("bitwise_and", x, y))
 
-@binary_ufunc(identity=0)
+@binary_ufunc(identity=0, reduce=reductions._reduce_bitwise_or)
 def bitwise_or(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the bitwise OR operation elementwise.
 
@@ -1279,7 +1279,7 @@ def bitwise_or(x: ArrayLike, y: ArrayLike, /) -> Array:
   """
   return lax.bitwise_or(*promote_args("bitwise_or", x, y))
 
-@binary_ufunc(identity=0)
+@binary_ufunc(identity=0, reduce=reductions._reduce_bitwise_xor)
 def bitwise_xor(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the bitwise XOR operation elementwise.
 
@@ -1793,16 +1793,7 @@ def spacing(x: ArrayLike, /) -> Array:
 
 
 # Logical ops
-def _logical_and_reduce(a: ArrayLike, axis: int = 0, dtype: DTypeLike | None = None,
-                        out: None = None, keepdims: bool = False, initial: ArrayLike | None = None,
-                        where: ArrayLike | None = None):
-  """Implementation of jnp.logical_and.reduce."""
-  if initial is not None:
-    raise ValueError("initial argument not supported in jnp.logical_and.reduce()")
-  result = reductions.all(a, axis=axis, out=out, keepdims=keepdims, where=where)
-  return result if dtype is None else result.astype(dtype)
-
-@binary_ufunc(identity=True, reduce=_logical_and_reduce)
+@binary_ufunc(identity=True, reduce=reductions._reduce_logical_and)
 def logical_and(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the logical AND operation elementwise.
 
@@ -1823,16 +1814,7 @@ def logical_and(x: ArrayLike, y: ArrayLike, /) -> Array:
   return lax.bitwise_and(*map(_to_bool, promote_args("logical_and", x, y)))
 
 
-def _logical_or_reduce(a: ArrayLike, axis: int = 0, dtype: DTypeLike | None = None,
-                       out: None = None, keepdims: bool = False, initial: ArrayLike | None = None,
-                       where: ArrayLike | None = None):
-  """Implementation of jnp.logical_or.reduce."""
-  if initial is not None:
-    raise ValueError("initial argument not supported in jnp.logical_or.reduce()")
-  result = reductions.any(a, axis=axis, out=out, keepdims=keepdims, where=where)
-  return result if dtype is None else result.astype(dtype)
-
-@binary_ufunc(identity=False, reduce=_logical_or_reduce)
+@binary_ufunc(identity=False, reduce=reductions._reduce_logical_or)
 def logical_or(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the logical OR operation elementwise.
 
@@ -1853,7 +1835,7 @@ def logical_or(x: ArrayLike, y: ArrayLike, /) -> Array:
   return lax.bitwise_or(*map(_to_bool, promote_args("logical_or", x, y)))
 
 
-@binary_ufunc(identity=False)
+@binary_ufunc(identity=False, reduce=reductions._reduce_logical_xor)
 def logical_xor(x: ArrayLike, y: ArrayLike, /) -> Array:
   """Compute the logical XOR operation elementwise.