[Pallas TPU] jnp.uint32 comparisons give the wrong result with large numbers

[ayaka14732]

Description

import functools
import jax
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
import jax.numpy as jnp

def main():
    @functools.partial(
        pl.pallas_call,
        out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM),
        out_shape=jax.ShapeDtypeStruct((2,), jnp.bool_),
    )
    def kernel(o_ref):
        o_ref[0] = jnp.uint32(5) <= jnp.uint32(4294967294)
        o_ref[1] = jnp.uint32(4294967292) > jnp.uint32(4)

    out = kernel()
    print(out)  # prints "[False False]"

if __name__ == '__main__':
    main()

Expected behaviour:

Outputs [True True] since 5 <= 4294967294 and 4294967292 > 4

Actual behaviour:

Outputs [False False]

Note that this only happens with large numbers. For example, 4294967294 exceeds the upper bound of jnp.int32, but does not exceed the upper bound of jnp.uint32. Therefore, this may be related to an unexpected cast from uint32 to int32 somewhere.

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.34.dev20240924+85a466d73
jaxlib: 0.4.33
numpy:  2.1.0
python: 3.12.4 (main, Jun  8 2024, 18:29:57) [GCC 11.4.0]
jax.devices (8 total, 8 local): [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0) TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0) ... TpuDevice(id=6, process_index=0, coords=(2,1,0), core_on_chip=0) TpuDevice(id=7, process_index=0, coords=(3,1,0), core_on_chip=0)]
process_count: 1
platform: uname_result(system='Linux', node='t1v-n-ab2ce832-w-0', release='5.19.0-1027-gcp', version='#29~22.04.1-Ubuntu SMP Thu Jun 22 05:13:17 UTC 2023', machine='x86_64')