def testNestedWhileWithDynamicUpdateSlice(self):
    num = 5

    def update_entry(arr, val, i, j):
        val = lax.reshape(val, [1, 1])
        return lax.dynamic_update_slice(arr, val, (i, j))

    def outer_loop(arr):

        def cond_fun(state):
            (i, num, _, _) = state
            return lax.lt(i, num)

        def body_fun(state):
            (i, num, arr, out) = state
            return (lax.add(i, 1), num, arr, inner_loop(i, arr, out))
        out = np.zeros(arr.shape, dtype=arr.dtype)
        init_val = (0, num, arr, out)
        (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
        return out

    def inner_loop(i, arr, out):

        def cond_fun(state):
            (i, j, _, _) = state
            return lax.le(j, i)

        def body_fun(state):
            (i, j, arr, out) = state
            arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
            arr_i_j = lax.dynamic_index_in_dim(arr_i, j, 0, False)
            out = update_entry(out, arr_i_j, i, j)
            return (i, lax.add(j, 1), arr, out)
        init_val = (i, 0, arr, out)
        (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
        return out
    cloop = api.jit(outer_loop)
    arr = npr.RandomState(0).randn(5, 5)
    self.assertAllClose(outer_loop(arr), np.tril(arr), check_dtypes=False)
    self.assertAllClose(cloop(arr), np.tril(arr), check_dtypes=False)
    self.assertAllClose(cloop(arr), np.tril(arr), check_dtypes=False)
def testNestedWhileWithDynamicUpdateSlice(self):
    num = 5

    def update_entry(arr, val, i, j):
        val = lax.reshape(val, [1, 1])
        return lax.dynamic_update_slice(arr, val, (i, j))

    def outer_loop(arr):

        def cond_fun(state):
            (i, num, _, _) = state
            return lax.lt(i, num)

        def body_fun(state):
            (i, num, arr, out) = state
            return (lax.add(i, 1), num, arr, inner_loop(i, arr, out))
        out = np.zeros(arr.shape, dtype=arr.dtype)
        init_val = (0, num, arr, out)
        (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
        return out

    def inner_loop(i, arr, out):

        def cond_fun(state):
            (i, j, _, _) = state
            return lax.le(j, i)

        def body_fun(state):
            (i, j, arr, out) = state
            arr_i = lax.dynamic_index_in_dim(arr, i, 0, False)
            arr_i_j = lax.dynamic_index_in_dim(arr_i, j, 0, False)
            out = update_entry(out, arr_i_j, i, j)
            return (i, lax.add(j, 1), arr, out)
        init_val = (i, 0, arr, out)
        (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
        return out
    cloop = api.jit(outer_loop)
    arr = npr.RandomState(0).randn(5, 5)
    self.assertAllClose(outer_loop(arr), np.tril(arr), check_dtypes=False)
    self.assertAllClose(cloop(arr), np.tril(arr), check_dtypes=False)
    self.assertAllClose(cloop(arr), np.tril(arr), check_dtypes=False)