Repository: google/jax

Issue:
lax.population_count does not support boolean dtype
Even though the use case is limited (`population_count(x) == x` if `x`'s dtype is `jax.numpy.bool_`), `lax.population_count_p` is defined as `population_count_p = standard_unop(_bool_or_int, 'population_count')`. However, this is what happens if `lax.population_count` is called with an argument with dtype `jax.numpy.bool_`:

```python
>>> lax.population_count(np.array([True, False], dtype=np.bool_))
Traceback (most recent call last):
  File "/jax/jax/interpreters/xla.py", line 311, in primitive_computation
    return c.build()
RuntimeError: Invalid argument: Expected an integral element type in argument to PopulationCount operation; got PRED.: 

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/jax/jax/lax/lax.py", line 296, in population_count
    return population_count_p.bind(x)
  File "/jax/jax/core.py", line 275, in bind
    return self.impl(*args, **kwargs)
  File "/jax/jax/interpreters/xla.py", line 224, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
  File "/jax/jax/interpreters/xla.py", line 257, in xla_primitive_callable
    built_c = primitive_computation(prim, AxisEnv(nreps), backend, tuple_args,
  File "/jax/jax/interpreters/xla.py", line 316, in primitive_computation
    raise RuntimeError(msg) from e
RuntimeError: Invalid argument: Expected an integral element type in argument to PopulationCount operation; got PRED.: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```

I think the simplest (and probably best) way to fix this is to remove booleans from the list of allowed types.


Discussion:
Good call! Either we could remove it from the allowed types, or cast to integer internally. I lean the same way you do, that removing the purported support for booleans is probably best.

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2242,7 +2242,7 @@ def _integer_pow_jvp(g, x, *, y):
 xor_p = standard_naryop([_bool_or_int, _bool_or_int], 'xor')
 ad.defjvp_zero(xor_p)
 
-population_count_p = standard_unop(_bool_or_int, 'population_count')
+population_count_p = standard_unop(_int, 'population_count')
 
 def _add_transpose(t, x, y):
   # The following linearity assertion is morally true, but because in some cases we
diff --git a/jax/lax_reference.py b/jax/lax_reference.py
--- a/jax/lax_reference.py
+++ b/jax/lax_reference.py
@@ -119,9 +119,14 @@ def rem(lhs, rhs):
 # TODO shift_right_logical
 
 def population_count(x):
+  assert np.issubdtype(x.dtype, np.integer)
   dtype = x.dtype
-  if x.dtype in (np.uint8, np.uint16):
-    x = x.astype(np.uint32)
+  iinfo = np.iinfo(x.dtype)
+  if np.iinfo(x.dtype).bits < 32:
+    assert iinfo.kind in ('i', 'u')
+    x = x.astype(np.uint32 if iinfo.kind == 'u' else np.int32)
+  if iinfo.kind == 'i':
+    x = x.view(f"uint{np.iinfo(x.dtype).bits}")
   assert x.dtype in (np.uint32, np.uint64)
   m = [
       0x5555555555555555,  # binary: 0101...


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -137,7 +137,7 @@ def op_record(op, nargs, dtypes, rng_factory, tol=None):
     op_record("bitwise_not", 1, bool_dtypes, jtu.rand_small),
     op_record("bitwise_or", 2, bool_dtypes, jtu.rand_small),
     op_record("bitwise_xor", 2, bool_dtypes, jtu.rand_small),
-    op_record("population_count", 1, uint_dtypes, jtu.rand_int),
+    op_record("population_count", 1, int_dtypes + uint_dtypes, jtu.rand_int),
 
     op_record("add", 2, default_dtypes + complex_dtypes, jtu.rand_small),
     op_record("sub", 2, default_dtypes + complex_dtypes, jtu.rand_small),
@@ -1786,6 +1786,12 @@ def test_reduction_with_repeated_axes_error(self):
     with self.assertRaisesRegex(ValueError, "duplicate value in 'axes' .*"):
       lax.reduce(np.arange(3), 0, lax.add, (0, 0))
 
+  def test_population_count_booleans_not_supported(self):
+    # https://github.com/google/jax/issues/3886
+    msg = "population_count does not accept dtype bool"
+    with self.assertRaisesRegex(TypeError, msg):
+      lax.population_count(True)
+
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):