Repository: jax-ml/jax
Issue category: Bug Report

Issue:
bincount rejects bool
### Description

[numpy.bincount](https://numpy.org/doc/stable/reference/generated/numpy.bincount.html) accepts bool, but [jax.numpy.bincount](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.bincount.html) does not:

```sh
$ py -c "import numpy as np; print(np.bincount(np.ones(3, bool)))"
[0 3]
$ py -c "from jax import numpy as jnp; print(jnp.bincount(jnp.ones(3, bool)))"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/carlos/venv/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py", line 3068, in bincount
    raise TypeError(f"x argument to bincount must have an integer type; got {_dtype(x)}")
TypeError: x argument to bincount must have an integer type; got bool
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.35
jaxlib: 0.4.34
numpy:  1.26.4
python: 3.12.7 (main, Oct  1 2024, 02:05:46) [Clang 15.0.0 (clang-1500.3.9.4)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='Carloss-MacBook-Pro-2.local', release='23.6.0', version='Darwin Kernel Version 23.6.0: Mon Jul 29 21:14:46 PDT 2024; root:xnu-10063.141.2~1/RELEASE_ARM64_T6031', machine='arm64')
```


Hints:

Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index d2e89833915d..b90004e19932 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3070,6 +3070,8 @@ def bincount(x: ArrayLike, weights: ArrayLike | None = None,
     Array([2, 1, 0, 1, 0], dtype=int32)
   """
   util.check_arraylike("bincount", x)
+  if _dtype(x) == bool:
+    x = lax.convert_element_type(x, 'int32')
   if not issubdtype(_dtype(x), integer):
     raise TypeError(f"x argument to bincount must have an integer type; got {_dtype(x)}")
   if ndim(x) != 1:
@@ -3080,7 +3082,7 @@ def bincount(x: ArrayLike, weights: ArrayLike | None = None,
     x_arr = core.concrete_or_error(asarray, x,
       "The error occurred because of argument 'x' of jnp.bincount. "
       "To avoid this error, pass a static `length` argument.")
-    length = max(minlength, x_arr.size and int(x_arr.max()) + 1)
+    length = max(minlength, x_arr.size and int(max(0, x_arr.max())) + 1)
   else:
     length = core.concrete_dim_or_error(length,
         "The error occurred because of argument 'length' of jnp.bincount.")


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 61baa7c97df4..7c2728af415e 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4905,7 +4905,7 @@ def testAtLeastNdLiterals(self, dtype, op):
 
   @jtu.sample_product(
     shape=[(0,), (5,), (10,)],
-    dtype=int_dtypes,
+    dtype=int_dtypes + bool_dtypes,
     weights=[True, False],
     minlength=[0, 20],
     length=[None, 8],