def testForiLoopTupleState(self):

    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)
        (_, total) = lax.fori_loop(0, lax.min(arr.shape[0], num), body_fun, init_val)
        return total
    cfun = api.jit(sum_first_n)
    x = npr.RandomState(0).randn(10).astype(jnp.float_)
    for num in [0, 5, 10, 15]:
        self.assertAllClose(sum_first_n(x, num), np.sum(x[:num]), check_dtypes=False)
        self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)
        self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)
def testForiLoopTupleState(self):

    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)
        (_, total) = lax.fori_loop(0, lax.min(arr.shape[0], num), body_fun, init_val)
        return total
    cfun = api.jit(sum_first_n)
    x = npr.RandomState(0).randn(10).astype(jnp.float_)
    for num in [0, 5, 10, 15]:
        self.assertAllClose(sum_first_n(x, num), np.sum(x[:num]), check_dtypes=False)
        self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)
        self.assertAllClose(cfun(x, num), np.sum(x[:num]), check_dtypes=False)