def sum_first_n(arr, num):

    def body_fun(i, state):
        (arr, total, _) = state
        arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
        return (arr, lax.add(total, arr_i), ())
    init_val = (arr, 0.0, ())
    (_, tot, _) = lax.fori_loop(0, lax.min(arr.shape[0], num), body_fun, init_val)
    return tot
def sum_first_n(arr, num):

    def body_fun(i, state):
        (arr, total, _) = state
        arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
        return (arr, lax.add(total, arr_i), ())
    init_val = (arr, 0.0, ())
    (_, tot, _) = lax.fori_loop(0, lax.min(arr.shape[0], num), body_fun, init_val)
    return tot