}
    return success();
  }
};

class ReplaceTpuPartitionedCallOpWithPartitionedCallOp
    : public OpRewritePattern<TF::TPUPartitionedCallOp> {
 public:
  using OpRewritePattern<TF::TPUPartitionedCallOp>::OpRewritePattern;

 private:
  LogicalResult matchAndRewrite(TF::TPUPartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {

}

bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
  bool has_tpu_partitioned_call = false;
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
    func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {
      has_tpu_partitioned_call = true;
    });
    if (has_tpu_partitioned_call) break;
  }
  return has_tpu_partitioned_call;
}

bool IsInferenceGraph(const Graph& graph,

  // a TF operation.
  bool isLegalToInline(Operation *call, Operation *callable,
                       bool wouldBeCloned) const final {
    // Skip inlining for TPUPartitionedCalls and RemoteCalls.
    if (isa<TPUPartitionedCallOp>(call)) return false;
    if (isa<RemoteCallOp>(call)) return false;
    // Maintain inlining for  `tf.function`s with jit_compile option.
    if (callable->hasAttr("tf._XlaMustCompile")) return true;

        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();
  if (!abstractOp) return false;
  return ops->count(abstractOp->getTypeID()) != 0;
}

  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto user : symbol_map.getUsers(func)) {
      // Populates work queue with func ops called from TPUPartionedCall.
      if (llvm::isa<TF::TPUPartitionedCallOp>(user)) {
        function_worklist.push(func);
      }
      // Populates caller to called func map.
      if (func::FuncOp caller = user->getParentOfType<func::FuncOp>()) {

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

template <typename CallOpClass>
static void SetPartitionCalleeFromCallable(CallOpClass op,

    if (!llvm::hasSingleElement(callers) &&
        !llvm::all_of(callers.drop_front(), [&](Operation* caller) {
          /// TODO(aminim): this is overly conservative as some operations
          /// (like TPUPartitionedCallOp) may have extra operands that aren't
          /// propagated to the callee.
          return isa<CallOpInterface>(caller) &&
                 std::equal(caller->getOperandTypes().begin(),