assert(while_wrapper && "While op is expected to be wrapped in a IslandOp");

  // Get the dummy constant.
  OpBuilder builder(while_wrapper);
  auto loc = NameLoc::get(
      builder.getStringAttr("chain_control_outputs@" + while_op.getBody()));
  IslandOp const_wrapper = GetDummyConstant(builder, const_type, loc);

  // Get new operand and result types.

  int index = *branch.getValues<int>().begin();
  if (index < 0 || index >= op.num_branches()) index = op.num_branches() - 1;

  auto func = mlir::cast<SymbolRefAttr>(op.getBranches()[index]);
  auto empty = rewriter.getStringAttr("");
  ReplaceTfOpWithNewOp<PartitionedCallOp>(
      rewriter, op, op.getResultTypes(), op.getOperands().drop_front(), func,
      /*config=*/empty, /*config_proto=*/empty, /*executor_type=*/empty);
  return success();

    ArrayAttr layout;  // filled in during the xla-adjust-layout pass
    auto data_and_token =
        rewriter.create<InfeedOp>(op.getLoc(), result_types, token,
                                  /*infeed_config=*/rewriter.getStringAttr(""),
                                  /*layout=*/layout);

    result_types.pop_back();  // remove the token type.

    if (op.get_XlaSharding().has_value()) {

  auto new_concat = builder.create<TFL::ConcatenationOp>(
      getLoc(), getType(), non_empty_values,
      builder.getIntegerAttr(builder.getIntegerType(32), getAxis()),
      builder.getStringAttr(getFusedActivationFunction()));
  return new_concat.getResult();
}

//===----------------------------------------------------------------------===//
// CustomOp