}

bool IsPartitionedOp(Operation* op) {
  static auto* ops = [] {
    llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
        new llvm::SmallDenseSet<mlir::TypeID, 32>{
            TypeID::get<TF::StatefulPartitionedCallOp>(),
            TypeID::get<TF::PartitionedCallOp>(),
            TypeID::get<TF::TPUPartitionedCallOp>(),
        };
    return ops_set;
  }();
  auto abstractOp = op->getRegisteredInfo();

    TypeID::get<TF::SqrtGradOp>(),
    TypeID::get<TF::SquaredDifferenceOp>(),
    TypeID::get<TF::SqueezeOp>(),
    TypeID::get<TF::StatelessParameterizedTruncatedNormalOp>(),
    TypeID::get<TF::StatefulPartitionedCallOp>(),
    TypeID::get<TF::StopGradientOp>(),
    TypeID::get<TF::StridedSliceOp>(),
    TypeID::get<TF::StridedSliceGradOp>(),
    TypeID::get<TF::SumOp>(),
    TypeID::get<TF::TanhGradOp>(),

// tf_device.launch op no longer exist after Outside Compilation is complete.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateVerifyNoOutsideCompilationMarkersPass();

// Create a pass that inlines the StatefulPartitionedCallOp op based in the
// parent region.
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaInlineDeviceOpsPass();

// Creates a pass that rewrites partitioned calls with `_xla_compile_device

                                      decomposed_partitioned_call_callees))) {
        return failure();
      }
    } else if (auto spcall =
                   llvm::dyn_cast<TF::StatefulPartitionedCallOp>(&op)) {
      if (failed(
              HandlePartitionedCallOp(spcall, spcall.func(), module, stats,
                                      decomposed_partitioned_call_callees))) {
        return failure();

  bool cond = cond_attr.getSplatValue<BoolAttr>().getValue();
  FlatSymbolRefAttr func =
      cond ? op.getThenBranchAttr() : op.getElseBranchAttr();

  // Replace IfOp with PartitionedCallOp or StatefulPartitionedCallOp.
  auto rewrite = [&](auto op_type) {
    auto empty = rewriter.getStringAttr("");
    ReplaceTfOpWithNewOp<typename decltype(op_type)::CallOp>(
        rewriter, op, op.getResultTypes(), op.getInput(), func,

  }

  return success();
}

LogicalResult PartitionedCallOp::verifySymbolUses(
    SymbolTableCollection &symbolTable) {
  return VerifyPartitionedCall(*this, symbolTable);
}
LogicalResult StatefulPartitionedCallOp::verifySymbolUses(
    SymbolTableCollection &symbolTable) {
  return VerifyPartitionedCall(*this, symbolTable);
}
LogicalResult TPUPartitionedCallOp::verifySymbolUses(
    SymbolTableCollection &symbolTable) {

          args_as_values.push_back(block_argument);
        }
        mlir::OpBuilder body_builder(&func.getBody());
        auto call = body_builder.create<mlir::TF::StatefulPartitionedCallOp>(
            func.getLoc(), orig_func.getFunctionType().getResults(),
            args_as_values,
            mlir::SymbolRefAttr::get(builder.getContext(), orig_func.getName()),