[JAX] Explicitly cast large integer constants to uint32.

We intend to add `jax.jit` decorators around a number of functions in the JAX standard library, including operators such as `&`.

A consequence of this is that JAX will attempt to cast Python integers (n.b. not NumPy scalars) to signed int32 or int64 types depending on the JAX x64 mode. The constant 2**32 - 1 is out of range of int32 (absent -x64 mode) and will produce an error under the new semantics. Instead, use a Numpy uint32 scalar.

PiperOrigin-RevId: 388950774

Author: hawkinsp

tensorflow_probability/python/internal/samplers.py

@@ -145,8 +145,8 @@ def fold_in(seed, salt):
   if JAX_MODE:
     from jax import random as jaxrand  # pylint: disable=g-import-not-at-top
     import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
-    return jaxrand.fold_in(seed,
-                           jnp.asarray(salt & (2**32 - 1), dtype=SEED_DTYPE))
+    return jaxrand.fold_in(
+        seed, jnp.asarray(salt & np.uint32(2**32 - 1), dtype=SEED_DTYPE))
   if isinstance(salt, (six.integer_types)):
     seed = tf.bitwise.bitwise_xor(
         seed, np.uint64([salt & (2**64 - 1)]).view(np.int32))