// while creating the ops so make sure there aren't any type changes between
    // the original op operands and the operands during the conversion.
    for (auto&& [old_val, new_val] : llvm::zip(op->getOperands(), operands)) {
      if (old_val.getType() != new_val.getType()) return failure();
    }

    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return failure();

  for (auto *op : custom_ops) {
    builder.setInsertionPoint(op);
    Location loc = op->getLoc();
    auto custom_op = builder.create<CustomTfOp>(loc, op->getResultTypes(),
                                                op->getOperands());
    Region region;
    Block *new_block = new Block;
    region.push_back(new_block);

    builder.setInsertionPointToEnd(&region.front());
    Operation *inner_op = builder.clone(*op);

        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();
    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
        return_op.getOperands()[0].getDefiningOp());
    if (!req_op) return failure();

    // Create a new func.call op with f32 output.
    auto new_call_op = call_op.clone();
    new_call_op->getResult(0).setType(

    // and mark the module as not valid for TFLite legalization.
    // Note: this might disable native variables in more than needed cases.
    // TODO(b/189370197): Enhance variable analysis.
    for (auto operand : op->getOperands()) {
      if (mlir::isa<TF::ResourceType>(
              getElementTypeOrSelf(operand.getType()))) {
        legalize_to_tfl = false;
        return WalkResult::interrupt();
      }
    }

  Operation* op = while_op.getOperation();
  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

    if (IsTrivialUnaryOperation(&cluster_op) ||
        HasOutsideCompilationAttribute(&cluster_op)) {
      auto walk_result = cluster_op.walk([&](Operation* op) {
        for (Value operand : op->getOperands()) {
          Operation* operand_op = GetOpOfValue(operand);
          if (head_outside_compiled_ops.count(operand_op)) continue;

          if (operand_op->getParentRegion() == cluster_region)

// 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(

    // which should be removed. So the replaced PartitionedCall op should keep
    // its original arguments except for the last element.
    SmallVector<Value> args = call_op.getOperands().drop_back();

    rewriter.replaceOpWithNewOp<TF::PartitionedCallOp>(
        call_op, float_func.getResultTypes(), args, f_attr);
    return success();
  }
};

template <typename OpT>
OpT AddOperandAndRewriteAs(Operation* op, Value operand, NamedAttrList attr,
                           OpBuilder* builder) {
  builder->setInsertionPoint(op);
  auto operands = llvm::to_vector<4>(op->getOperands());
  operands.push_back(operand);
  auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
                                     operands, attr.getAttrs());

      }
      // Map the block arguments to the yield results.
      for (auto iter_op : llvm::zip(for_op.getRegionIterArgs(),
                                    terminator_op->getOperands())) {
        mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
      }
      rewriter.eraseOp(terminator_op);
    }
    SmallVector<Value, 4> returned;
    for (Value arg : for_op.getRegionIterArgs()) {