def cond_fun(state):
    (arr, num, i, _) = state
    return lax.bitwise_and(lax.lt(i, num), lax.lt(i, arr.shape[0]))
def cond_fun(state):
    (arr, num, i, _) = state
    return lax.bitwise_and(lax.lt(i, num), lax.lt(i, arr.shape[0]))