def _while(*args: TfVal, cond_nconsts: int, cond_jaxpr: core.ClosedJaxpr, body_nconsts: int, body_jaxpr: core.ClosedJaxpr) -> Sequence[TfVal]:
    (cond_consts, body_consts, init_carry) = util.split_list(args, [cond_nconsts, body_nconsts])
    if cond_jaxpr.out_avals[0].shape:
        return _batched_cond_while(*args, cond_nconsts=cond_nconsts, cond_jaxpr=cond_jaxpr, body_nconsts=body_nconsts, body_jaxpr=body_jaxpr)

    def cond_tf_func(*args: TfVal) -> TfVal:
        (pred,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *args)
        return pred
    body_tf_func = functools.partial(_interpret_jaxpr, body_jaxpr, *body_consts)
    return tf.while_loop(cond_tf_func, body_tf_func, init_carry)
def _while(*args: TfVal, cond_nconsts: int, cond_jaxpr: core.ClosedJaxpr, body_nconsts: int, body_jaxpr: core.ClosedJaxpr) -> Sequence[TfVal]:
    (cond_consts, body_consts, init_carry) = util.split_list(args, [cond_nconsts, body_nconsts])
    if cond_jaxpr.out_avals[0].shape:
        return _batched_cond_while(*args, cond_nconsts=cond_nconsts, cond_jaxpr=cond_jaxpr, body_nconsts=body_nconsts, body_jaxpr=body_jaxpr)

    def cond_tf_func(*args: TfVal) -> TfVal:
        (pred,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *args)
        return pred
    body_tf_func = functools.partial(_interpret_jaxpr, body_jaxpr, *body_consts)
    return tf.while_loop(cond_tf_func, body_tf_func, init_carry)