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)
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)