// Verifies that the TF/XLA ops have all been lowered to MHLO.
std::unique_ptr<OperationPass<func::FuncOp>> CreateVerifyTFXLALegalizationPass(
    bool legalize_chlo = true);

// Transforms TFXLA Device specific ops into device independent ops.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTFXLADeviceSpecificTransformsPass(
    std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt);

    if (IsNonConstOp(op) && NotTFLQuantDequantizeOp(op) &&
        !IsTerminatorOp(op) &&
        !llvm::isa<func::ReturnOp, func::FuncOp, CallOpInterface>(op)) {
      SetTargetAnnotation(op, device_specs_flag_, &builder);
    }
  });
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
    llvm::ArrayRef<std::string> device_specs) {

// marked noinline should be specified by the `noinline-functions` option.
class MarkFunctionsNoinlinePass
    : public PassWrapper<MarkFunctionsNoinlinePass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MarkFunctionsNoinlinePass)

  explicit MarkFunctionsNoinlinePass()
      : MarkFunctionsNoinlinePass(

                                                 tensorflow::Session* session) {
  auto functions_range = module.getOps<func::FuncOp>();
  return mlir::failableParallelForEach(
      module.getContext(), functions_range.begin(), functions_range.end(),
      [&](func::FuncOp function) {
        return MarkInitializedVariablesInFunction(function, session);
      });
}

}  // namespace tf_saved_model

absl::Status AddSessionInitializerAndInlineCheckpoint(
    ModuleOp module, absl::string_view checkpoint_path) {
  // The main function should be the only public function.
  func::FuncOp main_func = nullptr;
  for (auto func : module.getOps<func::FuncOp>()) {
    if (func.isPublic()) {
      if (main_func != nullptr) {
        return absl::InternalError("Only one public function is allowed.");
      }
      main_func = func;

}

def BatchMatMulToEinsumPass : Pass<"tf-batch-matmul-to-tf-einsum", "mlir::func::FuncOp"> {
  let summary = "Replace TF BatchMatMul op by TF Einsum op.";
  let constructor = "TF::CreateBatchMatMulToEinsumPass()";
}

def CanonicalizeCompileAndReplicateAttributesPass : Pass<"tf-canonicalize-compile-and-replicate-attributes", "mlir::func::FuncOp"> {
  let summary = "Canonicalize compilation and replication attributes.";

  void TearDown() override {
    fused_lstm_func_.erase();
    fused_lstm_func_cifg_.erase();
    fused_ln_lstm_func_.erase();
    builder_.reset();
  }

  func::FuncOp fused_lstm_func_;
  func::FuncOp fused_lstm_func_cifg_;
  func::FuncOp fused_ln_lstm_func_;
  std::unique_ptr<mlir::MLIRContext> context_;
  std::unique_ptr<mlir::Builder> builder_;
};

TEST_F(LstmUtilsTest, ConvertLSTMCellSimple) {

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

namespace tfrt {
namespace fallback_async {

bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index);

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action);

void ForEachArgConsumedByFallback(
    mlir::ModuleOp module,

    module->emitError(
        absl::StrCat("failed to fetch device manager: ", status.message()));
    return failure();
  }

  SmallVector<func::FuncOp, 2> session_init_funcs =
      tf_saved_model::GetInitializerFunctions(module);
  func::FuncOp session_init_func =
      session_init_funcs.empty() ? nullptr : session_init_funcs[0];

  TF::ResourceAnalyzer analyzer(module, /*skip_session_init=*/true);

      total_element_count += input_type.getNumElements();
    }
  }
  return total_element_count;
}

struct GetOpCostPass
    : mlir::PassWrapper<GetOpCostPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetOpCostPass)

  llvm::StringRef getArgument() const final { return "tfl-get-op-cost"; }
  llvm::StringRef getDescription() const final {