def _batched_cond_while(*args: TfVal, cond_nconsts: int, cond_jaxpr: core.ClosedJaxpr, body_nconsts: int, body_jaxpr: core.ClosedJaxpr) -> Sequence[TfVal]:
    'Interprets a while_loop with a batched condition.\n\n  A batched while has a conditional that returns a tensor of booleans, and\n  a body that returns a list of tensors whose leading dimensions match those\n  of the conditional tensor.\n\n  We need to turn it into a while with scalar boolean conditional. We will\n  expand the loop carry to include a prefix with the current tensor boolean\n  condition. We prepend to the loop the first calculation of the tensor boolean\n  condition. The loop condition will use a "reduce_any" to calculate a scalar\n  boolean from the tensor boolean condition. The end of the loop body will\n  compute the new carry using a "tf.where", and we compute the new tensor\n  boolean condition.\n  '
    (cond_consts, body_consts, init_carry) = util.split_list(args, [cond_nconsts, body_nconsts])
    (init_pred_b,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *init_carry)
    assert (init_pred_b is not core.unit)

    def new_cond_tf_func(pred_b: TfVal, *carry: TfVal) -> TfVal:
        pred = tf.reduce_any(pred_b, axis=list(range(len(pred_b.shape))))
        return pred

    def new_body_tf_func(pred_b: TfVal, *carry: TfVal) -> Sequence[TfVal]:
        new_carry: Sequence[TfVal] = _interpret_jaxpr(body_jaxpr, *body_consts, *carry)

        def select_one_carry(new_c: TfVal, c: TfVal) -> TfVal:
            pred_b_bcast = _broadcast_in_dim(pred_b, shape=new_c.shape, broadcast_dimensions=list(range(len(pred_b.shape))))
            return tf.where(pred_b_bcast, new_c, c)
        selected_carry: Sequence[TfVal] = list(util.safe_map(select_one_carry, new_carry, carry))
        (next_pred_b,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *selected_carry)
        return (next_pred_b, *selected_carry)
    (_, *res_carry) = tf.while_loop(new_cond_tf_func, new_body_tf_func, (init_pred_b, *init_carry))
    return res_carry
Interprets a while_loop with a batched condition.

A batched while has a conditional that returns a tensor of booleans, and
a body that returns a list of tensors whose leading dimensions match those
of the conditional tensor.

We need to turn it into a while with scalar boolean conditional. We will
expand the loop carry to include a prefix with the current tensor boolean
condition. We prepend to the loop the first calculation of the tensor boolean
condition. The loop condition will use a "reduce_any" to calculate a scalar
boolean from the tensor boolean condition. The end of the loop body will
compute the new carry using a "tf.where", and we compute the new tensor
boolean condition.
def _batched_cond_while(*args: TfVal, cond_nconsts: int, cond_jaxpr: core.ClosedJaxpr, body_nconsts: int, body_jaxpr: core.ClosedJaxpr) -> Sequence[TfVal]:
    'Interprets a while_loop with a batched condition.\n\n  A batched while has a conditional that returns a tensor of booleans, and\n  a body that returns a list of tensors whose leading dimensions match those\n  of the conditional tensor.\n\n  We need to turn it into a while with scalar boolean conditional. We will\n  expand the loop carry to include a prefix with the current tensor boolean\n  condition. We prepend to the loop the first calculation of the tensor boolean\n  condition. The loop condition will use a "reduce_any" to calculate a scalar\n  boolean from the tensor boolean condition. The end of the loop body will\n  compute the new carry using a "tf.where", and we compute the new tensor\n  boolean condition.\n  '
    (cond_consts, body_consts, init_carry) = util.split_list(args, [cond_nconsts, body_nconsts])
    (init_pred_b,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *init_carry)
    assert (init_pred_b is not core.unit)

    def new_cond_tf_func(pred_b: TfVal, *carry: TfVal) -> TfVal:
        pred = tf.reduce_any(pred_b, axis=list(range(len(pred_b.shape))))
        return pred

    def new_body_tf_func(pred_b: TfVal, *carry: TfVal) -> Sequence[TfVal]:
        new_carry: Sequence[TfVal] = _interpret_jaxpr(body_jaxpr, *body_consts, *carry)

        def select_one_carry(new_c: TfVal, c: TfVal) -> TfVal:
            pred_b_bcast = _broadcast_in_dim(pred_b, shape=new_c.shape, broadcast_dimensions=list(range(len(pred_b.shape))))
            return tf.where(pred_b_bcast, new_c, c)
        selected_carry: Sequence[TfVal] = list(util.safe_map(select_one_carry, new_carry, carry))
        (next_pred_b,) = _interpret_jaxpr(cond_jaxpr, *cond_consts, *selected_carry)
        return (next_pred_b, *selected_carry)
    (_, *res_carry) = tf.while_loop(new_cond_tf_func, new_body_tf_func, (init_pred_b, *init_carry))
    return res_carry