namespace internal {

auto* has_tpu_partitioned_call_streamz =
    tensorflow::monitoring::Counter<1>::New(
        "/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call",
        "Whether the model has TPUPartitionedCallOp.",
        "has_tpu_partitioned_call");

namespace {

#define GEN_PASS_DEF_INFERENCEMETRICSPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h.inc"

    }
    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 {

func.func private @simple_callee(%arg0: tensor<2xi32>) -> tensor<2xi32> {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<2xi32>
  }
  func.return %graph#0 : tensor<2xi32>
}

// Test with a TPUPartitionedCallOp that is not inlined.
func.func @dont_inline_tpu_partitioned_call(%arg0: tensor<2xi32> {tf._user_specified_name = "inputs_0"}, %arg1: tensor<2xi32> {tf._user_specified_name = "inputs_1"}) -> tensor<2xi32> {

  %result = "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @simple_callee_with_noinline} : () -> tensor<2xi32>
  func.return %result : tensor<2xi32>
}

// Test that TPUPartitionedCallOp is not inlined.


// CHECK-LABEL: func @dont_inline_tpu_partitioned_call(
func.func @dont_inline_tpu_partitioned_call() -> tensor<2xi32> {
  // CHECK-NEXT: %[[ORDINAL:.*]] = "tf.TPUOrdinalSelector"

}

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(),