Fix jnp.diff for boolean inputs

Author: jakevdp

jax/_src/numpy/lax_numpy.py

@@ -1787,7 +1787,7 @@ def diff(a: ArrayLike, n: int = 1, axis: int = -1,
   slice1_tuple = tuple(slice1)
   slice2_tuple = tuple(slice2)
 
-  op = operator.not_equal if arr.dtype == np.bool_ else operator.sub
+  op = operator.ne if arr.dtype == np.bool_ else operator.sub
   for _ in range(n):
     arr = op(arr[slice1_tuple], arr[slice2_tuple])
 

tests/lax_numpy_test.py

@@ -2993,6 +2993,12 @@ def np_fun(x, n=n, axis=axis, prepend=prepend, append=append):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def testDiffBool(self):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng((10,), bool)]
+    self._CheckAgainstNumpy(np.diff, jnp.diff, args_maker, check_dtypes=False)
+    self._CompileAndCheck(jnp.diff, args_maker)
+
   def testDiffPrepoendScalar(self):
     # Regression test for https://github.com/jax-ml/jax/issues/19362
     x = jnp.arange(10)