CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

  op->replaceAllUsesWith(new_op.getResults());
  op->erase();
}

void LegalizeWhilePass::RunOnFunction(func::FuncOp func) {
  // Convert all TF WhileOps inside the function body to TFL While ops.
  func.getBody().walk([](TF::WhileOp while_op) { RunOnWhile(while_op); });
}

    auto index = indexed_assign.index();
    assign_op->moveAfter(new_replicate_op);
    assign_op->setOperand(
        1, new_replicate_op->getResult(old_num_results + num_replicas * index));
  }
  replicate_op->erase();
}

// Looks for AssignVariable ops from the end of the tf_device.replicate op. It
// returns all the last writes to replicate invariant resource variables

      ValueRange(resource_args), cluster_func_op.getFuncAttr());

  CopyDeviceAndUnderscoredAttributes(cluster_func_op, xla_launch_op);
  cluster_func_op.replaceAllUsesWith(xla_launch_op.getResults());
  cluster_func_op.erase();
}

void XlaRewritePass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symtab(module);
  OpBuilder builder(&getContext());
  module.walk([&](tf_device::ClusterFuncOp cluster_func_op) {

    DenseElementsAttr attr =
        DenseElementsAttr::get(type, b.getI32IntegerAttr(index));
    auto constant = b.create<arith::ConstantOp>(op.getLoc(), type, attr);
    op.replaceAllUsesWith(constant.getOperation());
    op.erase();
  });
}

// Creates an instance of the TensorFlow DeviceIndex selector pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDeviceIndexSelectorPass() {
  return std::make_unique<DeviceIndexSelector>();
}

            init_op.getLoc(), init_op.getTableHandle(), const_op.getResult(),
            init_op.getKeyIndex(), init_op.getValueIndex(),
            init_op.getVocabSize(), init_op.getDelimiter());
        init_op.erase();
      }
    }
    func.eraseArguments(args_to_erase);
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeAssetsPass(
    std::string saved_model_dir) {

    Operation &return_op = block.front();
    for (auto ret_mapping :
         llvm::zip(op->getResults(), return_op.getOperands())) {
      std::get<0>(ret_mapping).replaceAllUsesWith(std::get<1>(ret_mapping));
    }
    op->erase();
  });
}

}  // namespace

namespace TF {
std::unique_ptr<OperationPass<func::FuncOp>>
CreateMaterializePassthroughOpPass() {
  return std::make_unique<MaterializePassthroughOpPass>();
}

    TFL::FillOp fill = op_builder->create<TFL::FillOp>(
        const_op->getLoc(), splat_elements_attr.getType(), fill_shape,
        fill_value);
    const_op->replaceAllUsesWith(fill);
    const_op->erase();
  }
};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateUnfoldLargeSplatConstantPass() {
  return std::make_unique<UnfoldLargeSplatConstantPass>();
}

}  // namespace TFL

// results.
void RemoveIdentityFromRegion(Region& region) {
  region.walk([](Operation* op) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
      op->replaceAllUsesWith(op->getOperands());
      op->erase();
    }
  });
}

void TPUIdentityPruning::runOnOperation() {
  SmallVector<tf_device::ClusterOp, 4> clusters;
  getOperation().walk(
      [&](tf_device::ClusterOp cluster) { clusters.push_back(cluster); });

  it->second--;
  if (it->second != 0) {
    return;
  }
  tensor_usage_.erase(it);
  auto tensor_op_it = tensor_tape_.find(tensor_id);
  if (tensor_op_it == tensor_tape_.end()) {
    return;
  }
  const int64_t op_id = tensor_op_it->second;
  if (op_id == -1) {
    // Do not delete watched tensors.
    return;
  }
  tensor_tape_.erase(tensor_op_it);
  auto op_it = op_tape_.find(op_id);

    }
    args.append(new_args.begin(), new_args.end());
  } else {
    args.append(yield_op->operand_begin(), yield_op->operand_end());
  }
  b.create<func::ReturnOp>(yield_op->getLoc(), args);
  yield_op->erase();
  SymbolTable(region.getParentOfType<ModuleOp>()).insert(outlined_func);
  outlined_func.setPrivate();
  return outlined_func;
}

// Replace region with call to outline function.