def _fun(tup, _):
    (val, cond) = tup
    return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)
def _fun(tup, _):
    (val, cond) = tup
    return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)