//CHECK-NEXT: return %0#0 : tensor<448xf32>
//CHECK-NEXT:}

func.func @while(%init_i: tensor<i64>, %init_sum: tensor<i64>) -> tensor<i64>{
  %0, %1 = "vhlo.while_v1"(%init_i, %init_sum) ({
    ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
      %cond = "vhlo.compare_v1"(%arg0, %arg1) {
        comparison_direction = #vhlo<comparison_direction_v1 LT>,

  });
  module.walk([&](mlir::vhlo::WhileOpV1 while_op) {
    auto cond = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("cond").cast<mlir::FlatSymbolRefAttr>().getValue());
    InlineVhloOpRegion(while_op.getCond(), cond);
    while_op->removeAttr("cond");
    auto body = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("body").cast<mlir::FlatSymbolRefAttr>().getValue());

Michael Gester <******@****.***> 1654314705 -0700

Meghna Natraj <******@****.***> 1657816346 -0700

Meghna Natraj <******@****.***> 1657816346 -0700

Tres Popp <******@****.***> 1648205406 -0700

Michael Gester <******@****.***> 1654314705 -0700

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