Merge pull request #9901 from jakevdp:fix-int-x64

PiperOrigin-RevId: 434863868
google · Mar 15, 2022 · f47ce79 · f47ce79
2 parents 5354a01 + 36dabf1
commit f47ce79
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -4690,7 +4690,7 @@ def _unique(ar, axis, return_index=False, return_inverse=False, return_counts=Fa
         idx = idx.at[1:].set(where(idx[1:], idx[1:], mask.size))
       ret += (diff(idx),)
     elif ar.shape[axis]:
-      ret += (array([ar.shape[axis]], dtype=int_),)
+      ret += (array([ar.shape[axis]], dtype=dtypes.canonicalize_dtype(int_)),)
     else:
       ret += (empty(0, dtype=int),)
   if return_true_size: