"@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Support",
    ],
    alwayslink = 1,
)

# TODO(b/177376459): split tac_module and passes dependency to separate libraries.
cc_library(
    name = "target_aware_conversion",
    srcs = [
        "tac_module.cc",
        "transforms/compute_cost.cc",
        "transforms/fold_constants_to_subgraph.cc",
        "transforms/get_alternative_subgraph.cc",

        "  ).call,\n",
        "    with_gradient=False,\n",
        "    native_serialization=True,\n",
        "    native_serialization_platforms=('cpu',))\n",
        "\n",
        "tf_module = tf.Module()\n",
        "tf_module.f = tf.function(\n",
        "    jax_callable,\n",
        "    autograph=False,\n",
        "    input_signature=[\n",
        "        tf.TensorSpec(input_shape, jnp.float32, 'lhs_operand')\n",

  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  auto tf_module = mlir::parseSourceString<ModuleOp>(module_op_str, &context);
  if (!tf_module) {
    return absl::UnknownError("Failed to parse MLIR module");
  }

  auto mlir_file_path = ExportModule(*tf_module);
  if (!mlir_file_path.ok()) {
    return absl::UnknownError(

// module itself without inserting into the tf module.
FailureOr<StringAttr> RenameStablehloFunctions(
    MLIRContext *context, SymbolTableCollection &symbol_tables,
    ModuleOp tf_module, ModuleOp stablehlo_module) {
  SymbolTable &tf_symbol_table = symbol_tables.getSymbolTable(tf_module);
  // `stablehlo_module` is deleted right after the deserialization, so no need
  // to store its `SymbolTable` to `SymbolTableCollection`.

                    arith::ArithDialect, func::FuncDialect>();
  }

  explicit RaiseToTFOpsPass(std::optional<ModuleOp> tfr_module,
                            bool materialize_derived_attrs)
      : external_tfr_module_(tfr_module),
        materialize_derived_attrs_(materialize_derived_attrs) {}

  StringRef getArgument() const final { return "tfr-raise-to-tf"; }

  }
  return absl::OkStatus();
}

// Constructor of the decompose context.
TFRDecomposeContext::TFRDecomposeContext(mlir::ModuleOp tfr_module)
    : tfr_module_(tfr_module), pm_(tfr_module_.getContext()) {
  mlir::OpPassManager& func_pm = pm_.nest<mlir::func::FuncOp>();

  // Prepare the imported graph.

  os.flush();

  output->os() << result;
  output->keep();

  return absl::OkStatus();
}

tensorflow::Status ConvertTFToStableHLO(
    ModuleOp tf_module, const PassPipelineCLParser& pass_pipeline) {
  PassManager pm(tf_module.getContext());
  if (failed(applyPassManagerCLOptions(pm))) {
    return tensorflow::errors::Aborted(
        "Failed to apply MLIR pass manager CL options.");
  }

  func::FuncOp func =
      cluster_func->getParentOfType<ModuleOp>().lookupSymbol<func::FuncOp>(
          func_attr.getValue());

  std::string txt_module;
  if (failed(EncapsulateFuncAndSerialize(
          module_name.empty() ? "unknown_graph" : module_name.str(), func,
          &txt_module)))
    return nullptr;

  auto compilation_status_type =
      RankedTensorType::get({}, builder->getType<TF::StringType>());

    return;
  }

  // Print the module to a string before writing to the file.
  std::string txt_module;
  {
    llvm::raw_string_ostream os(txt_module);
    module.print(os);
  }

  status = file_writer->Append(txt_module);
  if (!status.ok()) {
    LOG(WARNING) << "error writing to file '" << prefix
                 << "': " << status.message();
    return;
  }

}

}  // namespace

// Creates an instance of the pass to decompose the TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
    std::optional<ModuleOp> tfr_module) {
  return std::make_unique<DecomposeTFOpsPass>(tfr_module);
}

static PassRegistration<DecomposeTFOpsPass> pass([] {
  return CreateDecomposeTFOpsPass();
});

}  // namespace TFR