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