Merge pull request #20258 from Gattocrucco:digitize-fix-corner-case

PiperOrigin-RevId: 616995961

Author: jax authors

jax/_src/numpy/lax_numpy.py

@@ -5373,7 +5373,7 @@ def digitize(x: ArrayLike, bins: ArrayLike, right: bool = False) -> Array:
   if bins_arr.ndim != 1:
     raise ValueError(f"digitize: bins must be a 1-dimensional array; got {bins=}")
   if bins_arr.shape[0] == 0:
-    return zeros(x, dtype=dtypes.canonicalize_dtype(int_))
+    return zeros_like(x, dtype=int32)
   side = 'right' if not right else 'left'
   return where(
     bins_arr[-1] >= bins_arr[0],

tests/lax_numpy_test.py

@@ -2597,7 +2597,7 @@ def testSearchsortedNans(self, dtype, side, method):
 
   @jtu.sample_product(
     xshape=[(20,), (5, 4)],
-    binshape=[(1,), (5,)],
+    binshape=[(0,), (1,), (5,)],
     right=[True, False],
     reverse=[True, False],
     dtype=default_dtypes,