// Whether the input file is an MLIR not tflite file.
    bool input_mlir = false;
  };

  explicit TfLiteImporter(const Options& options) : options_(options) {}

  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() override;

 private:
  Options options_;
  mlir::MLIRContext context_;
  llvm::SourceMgr source_mgr_;
  std::unique_ptr<mlir::SourceMgrDiagnosticHandler> source_mgr_handler_;
};

    : public impl::InferenceMetricsPassBase<InferenceMetricsPass> {
 public:
  void runOnOperation() override;
};

void InferenceMetricsPass::runOnOperation() {
  bool has_tpu_partitioned_call = false;
  ModuleOp module = getOperation();

  for (auto func_op : module.getOps<func::FuncOp>()) {
    func_op->walk(
        [&](TF::TPUPartitionedCallOp op) { has_tpu_partitioned_call = true; });

// Convert a TPU model to be compatible on CPU by rewriting/removing TPU ops.
class ConvertTpuModelToCpuPass
    : public PassWrapper<ConvertTpuModelToCpuPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertTpuModelToCpuPass)
  explicit ConvertTpuModelToCpuPass() = default;

  StringRef getArgument() const final {

    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
    if (!mlir_module_) {
      return absl::Status(
          absl::StatusCode::kNotFound,
          absl::StrCat("Could not find MLIR module at ", mlir_module_path));
    }
    return absl::OkStatus();
  }

  DialectRegistry registry_;
  MLIRContext context_;
  OwningOpRef<mlir::ModuleOp> mlir_module_;
};

// vhlo.op %1 : vhlo.tensor<...>
//   ==>
// vhlo.op %1 : tensor<...>
//
// TODO: There's likely a way to make MLIR manage the unrealized cast
// conversions using a specific rewriter.
LogicalResult ApplyTypeConverter(ModuleOp op, TypeConverter &converter) {
  IRRewriter rewriter(op->getContext());

  op->walk([&](Operation *op) {
    if (op->getDialect()->getNamespace() != "vhlo") return;

    // Convert operands

#include "tsl/platform/statusor.h"

namespace mlir {
namespace mhlo {
namespace test {

// Given a raw string, return a ModuleOp that can be used with the given
// MLIRContext.
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* mlir_context);

}  // namespace test
}  // namespace mhlo
}  // namespace mlir

#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

constexpr char kEntryFuncName[] = "main";

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {
  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);

  mlir::OwningOpRef<mlir::ModuleOp> module_ref;
  auto status =
      DeserializeMlirModule(str_attr.getValue().str(), context, &module_ref);
  if (!status.ok()) {
    LOG(ERROR) << status;
    return nullptr;
  }

  return module_ref;
}

static mlir::LogicalResult MlirModuleToSerializedMlirStringAttrTranslate(
    mlir::ModuleOp module_op, llvm::raw_ostream& output) {
  output << "\"";

// _xla_compile_device_type=TPU.
bool IsSupportedByReplicatedBridge(mlir::ModuleOp module);

// Check if an MLIR module contains TPUPartitionedCall op. If so, we define
// such graph as an inference graph. Otherwise, it is non inference graph.
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module);

// Check if a graph contains TPUPartitionedCall op, including its reachable

    TacFilter;
using ::third_party::tensorflow::compiler::mlir::lite::experimental::tac::
    TacFilters;

class TacFilterPass
    : public PassWrapper<TacFilterPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TacFilterPass)

  TacFilterPass() = default;
  TacFilterPass(const TacFilterPass& other) {
    this->tac_filters_ = other.tac_filters_;
  }