def sum_first_n(arr, num):

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

    def body_fun(state):
        (arr, num, i, total) = state
        arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
        return (arr, num, lax.add(i, 1), lax.add(total, arr_i))
    init_val = (arr, num, 0, 0.0)
    (_, _, _, total) = lax.while_loop(cond_fun, body_fun, init_val)
    return total
def sum_first_n(arr, num):

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

    def body_fun(state):
        (arr, num, i, total) = state
        arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
        return (arr, num, lax.add(i, 1), lax.add(total, arr_i))
    init_val = (arr, num, 0, 0.0)
    (_, _, _, total) = lax.while_loop(cond_fun, body_fun, init_val)
    return total