keras/src/backend/jax/random.py/randint
def randint(shape, minval, maxval, dtype="int32", seed=None):
    seed = jax_draw_seed(seed)
    return jax.random.randint(
        seed, shape=shape, dtype=dtype, minval=minval, maxval=maxval
    )