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,