def testLoopWithConjunctionCondition(self):

    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
    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 testLoopWithConjunctionCondition(self):

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