{"D", {PartialTensorShape({2, 3})}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoop) {
  // Graph:
  // x = array_ops.placeholder(dtypes.int32)
  // y = control_flow_ops.while_loop(lambda i: i < 10, lambda i: i + 1, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

## Known Caveats

*   TF Core:
    *   A longstanding bug in `tf.while_loop`, which caused it to execute
        sequentially, even when `parallel_iterations>1`, has now been fixed.
        However, the increased parallelism may result in increased memory use.
        Users who experience unwanted regressions should reset their
        `while_loop`'s `parallel_iterations` value to 1, which is consistent
        with prior behavior.

  for (auto extra_operand : extra_operands)
    new_types.push_back(extra_operand.getType());

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_types, operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());
  while_op.replaceAllUsesWith(

      return;
    }
    if (auto while_op = dyn_cast<TF::WhileOp>(op)) {
      for (auto callee : {while_op.cond_function(), while_op.body_function()}) {
        PropagatePotentiallyWrittenUpFromCallee(callee.getRegion(),
                                                while_op.getInput());
      }
      return;
    }
    if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {

  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    auto it = buffer_to_size->find(while_op.getOperand(i));
    if (it == buffer_to_size->end()) continue;
    new_while_operands.push_back(it->getSecond().size);
  }
  auto new_while = builder.create<TF::WhileRegionOp>(
      while_op.getLoc(), body_region.front().getTerminator()->getOperandTypes(),
      new_while_operands, while_op->getAttrs());

}

LogicalResult ConvertWhileOp(WhileOp while_op, bool allow_passthrough_args) {
  auto while_region = OpBuilder(while_op).create<TF::WhileRegionOp>(
      while_op.getLoc(), while_op.getResultTypes(), while_op.getInput(),
      while_op.getParallelIterations(), while_op.getIsStateless(),
      while_op.getShapeInvariant());
  CopyDeviceAndUnderscoredAttributes(while_op, while_region);

  YieldOp cond_yield =

      id_map[i] = j++;
      new_operands.push_back(while_op.getOperand(i));
      new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {

      read_var->erase();
    } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) {
      GetUpdatedWhileOp<TF::WhileOp>(
          while_op, while_op.cond_function().getArgumentTypes(), args_to_erase);
      while_op->erase();
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(user_op)) {
      auto new_while_op = GetUpdatedWhileOp(
          while_op, while_op.getCond().getArgumentTypes(), args_to_erase);

  // Create the new while op.
  auto new_while_operands = llvm::to_vector<8>(while_op.getOperands());
  OpBuilder builder(while_op);
  assert(while_op.getNumOperands() == while_op.getNumResults());
  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    auto it = data_var_to_size_var.find(while_op.getOperand(i));
    if (it == data_var_to_size_var.end()) continue;

llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4>
AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping(
    TF::WhileRegionOp while_op, tf_device::ReplicateOp replicate,
    TF::TPUExecuteAndUpdateVariablesOp execute,
    tf_device::LaunchOp compile_launch) {
  Region& body = while_op.getBody();
  Region& cond = while_op.getCond();

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4> mapping;