Merge pull request #17962 from jakevdp:fix-bitwise-count

PiperOrigin-RevId: 571398043
jax-ml · Oct 6, 2023 · 6681e64 · 6681e64
2 parents fae53d9 + cd18f8e
commit 6681e64
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
@@ -201,7 +201,7 @@ def arccosh(x: ArrayLike, /) -> Array:
 @jit
 def bitwise_count(x: ArrayLike, /) -> Array:
   # Following numpy we take the absolute value and return uint8.
-  return lax.population_count(lax.abs(x)).astype('uint8')
+  return lax.population_count(abs(x)).astype('uint8')
 
 @_wraps(np.right_shift, module='numpy')
 @partial(jit, inline=True)

diff --git a/tests/lax_numpy_operators_test.py b/tests/lax_numpy_operators_test.py
@@ -581,7 +581,7 @@ def testBitwiseOp(self, name, rng_factory, shapes, dtypes):
 
   @jtu.sample_product(
     shape=array_shapes,
-    dtype=int_dtypes,
+    dtype=int_dtypes + unsigned_dtypes,
   )
   def testBitwiseCount(self, shape, dtype):
     # np.bitwise_count added after numpy 1.26, but