# Buggy code
def population_count(x):
    dtype = x.dtype
    if x.dtype in (np.uint8, np.uint16):
        x = x.astype(np.uint32)
    assert x.dtype in (np.uint32, np.uint64)
    m = [
        0x5555555555555555,  # binary: 0101...
        0x3333333333333333,  # binary: 00110011..
        0x0F0F0F0F0F0F0F0F,  # binary:  4 zeros,  4 ones ...
        0x00FF00FF00FF00FF,  # binary:  8 zeros,  8 ones ...
        0x0000FFFF0000FFFF,  # binary: 16 zeros, 16 ones ...
        0x00000000FFFFFFFF,  # binary: 32 zeros, 32 ones
    ]

    if x.dtype == np.uint32:
        m = list(map(np.uint32, m[:-1]))
    else:
        m = list(map(np.uint64, m))

    x = (x & m[0]) + ((x >> 1) & m[0])  # put count of each  2 bits into those  2 bits
    x = (x & m[1]) + ((x >> 2) & m[1])  # put count of each  4 bits into those  4 bits
    x = (x & m[2]) + ((x >> 4) & m[2])  # put count of each  8 bits into those  8 bits
    x = (x & m[3]) + ((x >> 8) & m[3])  # put count of each 16 bits into those 16 bits
    x = (x & m[4]) + ((x >> 16) & m[4])  # put count of each 32 bits into those 32 bits
    if x.dtype == np.uint64:
        x = (x & m[5]) + (
            (x >> 32) & m[5]
        )  # put count of each 64 bits into those 64 bits
    return x.astype(dtype)


# Traceback (RuntimeError)
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

# Fixed code
def population_count(x):
    assert np.issubdtype(x.dtype, np.integer)
    dtype = x.dtype
    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...
        0x3333333333333333,  # binary: 00110011..
        0x0F0F0F0F0F0F0F0F,  # binary:  4 zeros,  4 ones ...
        0x00FF00FF00FF00FF,  # binary:  8 zeros,  8 ones ...
        0x0000FFFF0000FFFF,  # binary: 16 zeros, 16 ones ...
        0x00000000FFFFFFFF,  # binary: 32 zeros, 32 ones
    ]

    if x.dtype == np.uint32:
        m = list(map(np.uint32, m[:-1]))
    else:
        m = list(map(np.uint64, m))

    x = (x & m[0]) + ((x >> 1) & m[0])  # put count of each  2 bits into those  2 bits
    x = (x & m[1]) + ((x >> 2) & m[1])  # put count of each  4 bits into those  4 bits
    x = (x & m[2]) + ((x >> 4) & m[2])  # put count of each  8 bits into those  8 bits
    x = (x & m[3]) + ((x >> 8) & m[3])  # put count of each 16 bits into those 16 bits
    x = (x & m[4]) + ((x >> 16) & m[4])  # put count of each 32 bits into those 32 bits
    if x.dtype == np.uint64:
        x = (x & m[5]) + (
            (x >> 32) & m[5]
        )  # put count of each 64 bits into those 64 bits
    return x.astype(dtype)