Inconsistent edge case handling in jax.random.categorical across devices
When provided NaN probabilities, `jax.random.categorical` returns different results depending on the device. 

```
# GPU
> jax.random.categorical(rng_key, jnp.array([0, jnp.nan]))
DeviceArray(2147483647, dtype=int32)

# CPU
> jax.random.categorical(rng_key, jnp.array([0, jnp.nan]))
DeviceArray(1, dtype=int32)
```

Arguably both are not ideal (being valid values, bugs can go unnoticed - especially in the CPU case where 1 is a perfectly valid value coming from this function) but the inconsistency in particular can make debugging very tricky and should be fixed.
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 05fb817f7..8e1ff56f7 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5382,7 +5382,8 @@ def _argminmax_translation_rule(value_comparator, identity,
   x_index = xb.parameter(subc, 1, index_shape)
   y_value = xb.parameter(subc, 2, value_shape)
   y_index = xb.parameter(subc, 3, index_shape)
-  which_value = value_comparator(x_value, y_value)
+  which_value = xops.Or(value_comparator(x_value, y_value),
+                        xops.Ne(x_value, x_value))
   which_index = xops.Or(which_value, xops.And(xops.Eq(x_value, y_value),
                                               xops.Lt(x_index, y_index)))
   xops.Tuple(subc, [xops.Select(which_value, x_value, y_value),
@@ -5402,8 +5403,8 @@ def _argminmax_gpu_translation_rule(op, a, *, axes, index_dtype):
   idxs = tie_in(a, broadcasted_iota(index_dtype, a.shape, axis))
   maxval = np.array(dtypes.iinfo(index_dtype).max, dtype=index_dtype)
   maxval = broadcast(tie_in(a, maxval), a.shape)
-  mask_idxs = select(eq(a, expand_dims(op(a, (axis,)), (axis,))), idxs,
-                     maxval)
+  maxvals = expand_dims(op(a, (axis,)), (axis,))
+  mask_idxs = select(eq(a, maxvals) | ne(a, a), idxs, maxval)
   return _reduce_min(mask_idxs, (axis,))
 
 _argmin_translation_rule = partial(_argminmax_translation_rule, xops.Lt,
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 9ca26b04e..15817ae75 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2560,6 +2560,10 @@ class LazyConstantTest(jtu.JaxTestCase):
     x_out_jit = api.jit(op)(x_in)
     self.assertEqual(dtypes.is_weakly_typed(x_out_jit), False)
 
+  def testArgMaxOfNanChoosesNaN(self):
+    self.assertEqual(lax.argmax(np.array([0., np.nan]), axis=0,
+                                index_dtype=np.int32), 1)
+
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name": "_{}".format(rec.op),
          "op_name": rec.op, "rec_dtypes": rec.dtypes}
['jax/_src/lax/lax.py', 'tests/lax_test.py']