class TacImporter {
 public:
  virtual ~TacImporter() = default;

  // Imports and returns the Module for the imported program.
  virtual absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() = 0;
};

// Interface for exporting a module.
// Users should implement the interface for exporting the result from TAC
// in their preferred way.

mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module, std::function<WalkResult(Operation*, tf_device::ClusterOp,
                                              std::optional<std::string>)>
                         callback);

// Like above, except TPU clusters are not required to have a host device, and
// no host device is passed to `callback`.
mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module,

static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT";

namespace {
inline absl::string_view StringRefToView(llvm::StringRef ref) {
  return {ref.data(), ref.size()};
}
}  // namespace

static OwningOpRef<mlir::ModuleOp> GraphdefToMlirTranslateFunction(
    llvm::StringRef input, MLIRContext* context) {
  tensorflow::GraphdefToMlirOptions options{
      debug_info_file,        xla_compile_device_type,

    };
    return dump_attributes;
  }

  StringAttr DuplicateFunction(Operation *op,
                               const FlatSymbolRefAttr &f_attr) const {
    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);

    const func::FuncOp ref_func =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(f_attr.getValue()));

TEST(BridgeLoggerFilters, TestPassFilter) {
  mlir::DialectRegistry mlir_registry;
  mlir::RegisterAllTensorFlowDialects(mlir_registry);
  mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_with_add;
  TF_ASSERT_OK(DeserializeMlirModule(module_with_add, &mlir_context,
                                     &mlir_module_with_add));

  std::unique_ptr<mlir::Pass> partitioning_pass =

          absl::StrCat("Failed to parse input MLIR model: ", error_message));
    }

    llvm::SourceMgr source_mgr;
    source_mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    auto module = parseSourceFile<ModuleOp>(source_mgr, context);
    if (module->getOperation() == nullptr) {
      return absl::AbortedError("Failed to parse input MLIR model.");
    }

    return quant::stablehlo::ImportedMlirModuleOp(std::move(module), nullptr);

using ::tensorflow::quantization::RunPasses;

PostCalibrationComponent::PostCalibrationComponent(
    absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> PostCalibrationComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(
      kName, /*add_passes_func=*/
      [&config](PassManager& pm) {

void XlaInlineDeviceOpsPass::runOnOperation() {
  ModuleOp module = getOperation();

  llvm::SmallVector<tf_device::ClusterOp, 4> ops;
  module.walk(
      [&](tf_device::ClusterOp cluster_op) { ops.push_back(cluster_op); });

  for (auto cluster_op : ops) {
    InlineDeviceOp(cluster_op);
  }
}

}  // namespace

namespace TFDevice {
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaInlineDeviceOpsPass() {

    const llvm::SmallVector<int64_t>& expected_packed_shape,
    const llvm::SmallVector<int8_t>& expected_packed_values) {
  MLIRContext context;
  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
  OpBuilder builder(&module->getBodyRegion());
  context.loadDialect<TF::TensorFlowDialect>();

  Value value = CreateConstValue<int8_t>(builder, module->getLoc(),

namespace tensorflow {
namespace {

constexpr char kCostAttrName[] = "_tfrt_cost";
constexpr char kOpKeyAttrName[] = "op_key";

absl::flat_hash_map<int64_t, uint64_t> GetOpCostMap(mlir::ModuleOp op) {
  absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {
    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
    if (!cost_attr) return;