// 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); });

    }
    custom_op->setAttrs(inner_op->getAttrs());
    builder.create<YieldOp>(loc, inner_op->getResults());
    custom_op.getBody().takeBody(region);

    op->replaceAllUsesWith(custom_op);
    op->erase();
  }
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass() {
  return std::make_unique<RaiseCustomOpsPass>();
}

    op.replaceAllUsesWith(existing_resource->getParentOp()->getResults());
    ops_to_remove.push_back(&op);
  }

  // Remove op after the loop to avoid crash.
  for (Operation* op : ops_to_remove) {
    op->erase();
  }
}

static PassRegistration<MergeDuplicateResourceOpsPass> pass{};

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>
CreateMergeDuplicateResourceOpsPass() {

                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return absl::OkStatus();
}

  // Wrap the cluster op into a Launch op
  auto launch_op = tensorflow::WrapOpInLaunch(&builder, loc, cluster, device);

  EXPECT_TRUE(llvm::isa<mlir::tf_device::LaunchOp>(launch_op));
  launch_op->erase();
}

}  // namespace

  auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
                                     operands, attr.getAttrs());
  op->replaceAllUsesWith(new_op.getOperation()->getResults());
  op->erase();
  return new_op;
}

// Returns success if the function has at most one op of the template type and
// assigns it to `result`, if present. If there are multiple such ops, returns
// failure.
template <typename OpT>

          LLVM_DEBUG(llvm::dbgs() << "Re-use sunk constant " << use->get()
                                  << "\n     in " << use->get() << "\n");
          if (constant.use_empty()) const_op.erase();
          return;
        }
        if (constant.hasOneUse()) {
          LLVM_DEBUG(llvm::dbgs() << "Moved constant " << constant << "\n");
          const_op.getOperation()->moveBefore(&body.begin()->front());

    if (!best) {
      assert(ready.empty());
      return result;  // happens for unused results for ops in the todo list
    }

    // Consider this operation emitted, and make its results available.
    ready.erase(std::find(ready.begin(), ready.end(), best));
    previous_op = best;
    for (Value result : best->getResults()) {
      todo.push(result);
    }
    for (Operation* successor : ctrlSuccessors(best)) {

    for (auto [result, new_result] :
         llvm::zip_equal(op.getResults(), new_op.getResults())) {
      new_op_mapper.map(result, new_result);
    }
  }

  call_op.erase();
}

void UnwrapXlaCallModuleOpPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  SymbolTable symbol_table(module_op);

  for (auto func_op : module_op.getOps<func::FuncOp>()) {

                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return absl::OkStatus();
}