#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/managed_stack_trace.h"

namespace mlir {

StatusScopedDiagnosticHandler::StatusScopedDiagnosticHandler(
    MLIRContext* context, bool propagate, bool filter_stack)
    : BaseScopedDiagnosticHandler(context, propagate, filter_stack) {
  if (filter_stack) {

namespace {

using testing::HasSubstr;

TEST(ErrorUtilTest, StatusScopedDiagnosticHandler) {
  MLIRContext context;
  auto id = StringAttr::get(&context, "//tensorflow/python/test.py");
  auto loc = FileLineColLoc::get(&context, id, 0, 0);

  // Test OK without diagnostic gets passed through.
  {
    TF_ASSERT_OK(
        StatusScopedDiagnosticHandler(&context).Combine(absl::OkStatus()));
  }

      TF_SetStatus(status, TF_INVALID_ARGUMENT,
                   ("Invalid pass_pipeline: " + error_stream.str()).c_str());
      return "// error";
    }

    mlir::StatusScopedDiagnosticHandler statusHandler(module.getContext());
    if (failed(pm.run(module))) {
      tsl::Set_TF_Status_from_Status(status, statusHandler.ConsumeStatus());
      return "// error";
    }
  }

namespace tensorflow {
namespace quantization {

absl::Status RunPassesOnModuleOp(
    std::optional<absl::string_view> mlir_dump_file_name,
    mlir::PassManager& pass_manager, mlir::ModuleOp module_op) {
  mlir::StatusScopedDiagnosticHandler statusHandler(module_op.getContext(),
                                                    /*propagate=*/true);

  absl::StatusOr<std::unique_ptr<llvm::raw_ostream>> dump_file;
  if (mlir_dump_file_name) {

  TF_RET_CHECK(mlir_module) << "unexpected null MLIR module pointer";

  // Make sure we catch any error reported by MLIR and forward it to the TF
  // error reporting system.
  mlir::StatusScopedDiagnosticHandler error_handler(mlir_context);

  // Parse the module.
  *mlir_module = mlir::parseSourceString<mlir::ModuleOp>(serialized_mlir_module,
                                                         mlir_context);

absl::Status RunPasses(const absl::string_view name, FuncT add_passes_func,
                       mlir::MLIRContext& ctx, mlir::ModuleOp module_op) {
  mlir::PassManager pm{&ctx};
  add_passes_func(pm);

  mlir::StatusScopedDiagnosticHandler diagnostic_handler{&ctx};
  TF_RETURN_IF_ERROR(MaybeEnableIrPrinting(pm, name));

  if (failed(pm.run(module_op))) {
    return absl::InternalError(
        absl::StrFormat("Failed to run pass: %s. %s", name,

  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> module;
  if (flib_def) {
    flib_def = &graph.flib_def();
  }
  auto convert = [&]() -> Status {
    mlir::StatusScopedDiagnosticHandler status_handler(&context);
    // TODO(jpienaar): Both the graph debug info and import config should be
    // specifiable.
    GraphDebugInfo debug_info;
    switch (config.dialect) {

  // Prepare IR for exporting.
  pm.addPass(CreateBreakUpIslandsPass());

  // In case of failure, the `diag_handler` converts MLIR errors emitted to the
  // MLIRContext into a tensorflow::Status.
  StatusScopedDiagnosticHandler diag_handler(module.getContext());
  LogicalResult result = pm.run(module);
  (void)result;
  return diag_handler.ConsumeStatus();
}

}  // namespace TF

  print_flags.elideLargeElementsAttrs();

  if (VLOG_IS_ON(1)) {
    VLOG(1) << "Input TF Executor dialect:";
    DumpMlirOpToFile("tf_to_tfrt_tf_executor_dialect", module);
  }

  mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());

  // Lower MLIR TF Dialect to MLIR TFRT CoreRT dialect.
  mlir::PassManager pm(module.getContext());
  tensorflow::applyTensorflowAndCLOptions(pm);

namespace mlir {
namespace lite {

absl::Status SparsifyModel(const tflite::ModelT& input_model,
                           flatbuffers::FlatBufferBuilder* builder) {
  MLIRContext context;
  StatusScopedDiagnosticHandler statusHandler(&context,
                                              /*propagate=*/true);

  // Import input_model to a MLIR module
  flatbuffers::FlatBufferBuilder input_builder;