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), ())
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), ())