Repository: google/jax

Issue:
Unexpected behavior of `segment_max` with boolean arrays
## Problem

`segment_max` applied to boolean arrays seems to generally return `True` for all segments. This is in contrast to the behavior of `segment_min` as well as `min` and `max`. 

## Example

```python
import jax.numpy as jnp
from jax.ops import segment_max

# max over boolean array
jnp.array([False, False]).max()
# DeviceArray(False, dtype=bool); As expected

# min over boolean array
jnp.array([True, False]).min()
# DeviceArray(False, dtype=bool); As expected

# segment_max over boolean array
segment_max(
    jnp.array([True, False, False, False]),
    jnp.array([0, 0, 1, 1])
)
# DeviceArray([ True,  True], dtype=bool); Expected result was [True, False]

# segment_min over boolean array
segment_min(
    jnp.array([True, True, False, False]),
    jnp.array([0, 0, 1, 1])
)
# DeviceArray([ True, False], dtype=bool); As expected

```
## Package versions

jax: 0.3.0
jaxlib: 0.3.0


Discussion:
Thanks for the report - I think it's an issue with the identity for this computation. I'll work on a fix.

Resolving patch:
diff --git a/jax/_src/ops/scatter.py b/jax/_src/ops/scatter.py
--- a/jax/_src/ops/scatter.py
+++ b/jax/_src/ops/scatter.py
@@ -115,7 +115,6 @@ def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
 
 
 
-
 def _get_identity(op, dtype):
   """Get an appropriate identity for a given operation in a given dtype."""
   if op is lax.scatter_add:
@@ -123,11 +122,15 @@ def _get_identity(op, dtype):
   elif op is lax.scatter_mul:
     return 1
   elif op is lax.scatter_min:
-    if jnp.issubdtype(dtype, jnp.integer):
+    if dtype == dtypes.bool_:
+      return True
+    elif jnp.issubdtype(dtype, jnp.integer):
       return jnp.iinfo(dtype).max
     return float('inf')
   elif op is lax.scatter_max:
-    if jnp.issubdtype(dtype, jnp.integer):
+    if dtype == dtypes.bool_:
+      return False
+    elif jnp.issubdtype(dtype, jnp.integer):
       return jnp.iinfo(dtype).min
     return -float('inf')
   else:


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -1229,6 +1229,49 @@ def fn(data, segment_ids):
     self.assertAllClose(grad, np.array([0., 0.], np.float32))
 
 
+  @parameterized.named_parameters(itertools.chain.from_iterable(
+      jtu.cases_from_list({
+        "testcase_name": "_{}_{}_num_segments={}_bucket_size={}".format(
+          jtu.format_shape_dtype_string(shape, dtype),
+          reducer.__name__, num_segments, bucket_size),
+        "dtype": dtype, "shape": shape,
+        "reducer": reducer, "op": op, "identity": identity,
+        "num_segments": num_segments, "bucket_size": bucket_size}
+      for dtype in [np.bool_]
+      for shape in [(8,), (7, 4), (6, 4, 2)]
+      for bucket_size in [None, 2]
+      for num_segments in [None, 1, 3])
+    for reducer, op, identity in [
+      (ops.segment_min, np.minimum, True),
+      (ops.segment_max, np.maximum, False),
+    ]))
+  def testSegmentReduceBoolean(self, shape, dtype, reducer, op, identity, num_segments, bucket_size):
+    rng = jtu.rand_default(self.rng())
+    idx_rng = jtu.rand_int(self.rng(), low=-2, high=3)
+    args_maker = lambda: [rng(shape, dtype), idx_rng(shape[:1], jnp.int32)]
+
+    if np.issubdtype(dtype, np.integer):
+      if np.isposinf(identity):
+        identity = np.iinfo(dtype).max
+      elif np.isneginf(identity):
+        identity = np.iinfo(dtype).min
+
+    jnp_fun = lambda data, segment_ids: reducer(
+      data, segment_ids, num_segments=num_segments, bucket_size=bucket_size)
+
+    def np_fun(data, segment_ids):
+      size = num_segments if num_segments is not None else (segment_ids.max() + 1)
+      out = np.full((size,) + shape[1:], identity, dtype)
+      for i, val in zip(segment_ids, data):
+        if 0 <= i < size:
+          out[i] = op(out[i], val).astype(dtype)
+      return out
+
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+    if num_segments is not None:
+      self._CompileAndCheck(jnp_fun, args_maker)
+
+
   @parameterized.named_parameters(itertools.chain.from_iterable(
       jtu.cases_from_list({
         "testcase_name": "_{}_{}_num_segments={}_bucket_size={}".format(