def testWhileWithManyArgs(self):
    nargs = 256

    def loop_cond(state):
        return lax.lt(state[0], 2)

    def loop_body(state):
        return tuple((lax.add(s, 1) for s in state))
    _ = lax.while_loop(loop_cond, loop_body, ((0,) * nargs))
def testWhileWithManyArgs(self):
    nargs = 256

    def loop_cond(state):
        return lax.lt(state[0], 2)

    def loop_body(state):
        return tuple((lax.add(s, 1) for s in state))
    _ = lax.while_loop(loop_cond, loop_body, ((0,) * nargs))