Fixed issue in jnp.argpartition and jnp.partition with unsigned integers with zeros

[vfdev-5]

Fixes #22137

Description:

• Fixed issue in jnp.argpartition and jnp.partition with unsigned integers with zeros

Thanks @pearu for the trick with -lax.top_k(-(arr + 1), kth + 1)[0] - 1

Details on the problem with unsigned integers:

import jax.numpy as jnp
from jax import lax

a = jnp.array([0, 1, 2, 254, 255], dtype=jnp.uint8)
print(-a)   # [  0 255 254   2   1]

# On main now:
print(lax.top_k(-a, 3)[0])  # [255 254   2]
# but it should be [0  255 254]
print(-lax.top_k(-a, 3)[0])  # [  1   2 254]
# but it should be [0  1   2]

# On this PR using the trick suggested by Pearu
print(lax.top_k(-(a + 1), 3)[0])  # [255 254 253]
print(-lax.top_k(-(a + 1), 3)[0] - 1)  # [0 1 2]

[vfdev-5]

I refactored the assertion part between testPartition and testArgpartition functions and reused this method in added tests. If there can be another way of doing the tests I'm happy to implement it

[jakevdp]

Nice!