"Resource ops decomposition did not converge";

// Decomposing resource ops should not take more than a few iterations (2-3) to
// converge as only a few patterns create new resource ops that can be further
// decomposed. The rest of the iterations are enough to clean up any dead ops
// created by decomposition.
constexpr int kMaxIterations = 20;

// Populates `reachable_functions` with all functions that can be reached from
// device cluster ops.

        pm.addPass(mlir::quant::CreateUnfreezeConstantsPass());
      },
      ctx, module_op));

  if (const absl::Status create_dir_status =
          Env::Default()->CreateDir(std::string(checkpoint_dir));
      !create_dir_status.ok()) {
    LOG(ERROR) << "Failed to create checkpoint directory at: "
               << checkpoint_dir;
    return create_dir_status;
  }

  TF_ASSIGN_OR_RETURN(const auto unused_variable_names,

                                            &capture_handle));
    captures.push_back(capture_handle);
  }
  // TODO(bmzhao): Create Metadata here
  return TFConcreteFunction::Create(/*function_def=*/builder.fdef,
                                    /*captures=*/std::move(captures),
                                    /*metadata=*/{},
                                    /*ctx=*/ctx,

      // Create new ConstantOp-ConvertOp-Operation sequences. At this moment,
      // old ConstantOp is guaranteed to have one F32->F16 convert op regardless
      // of its number of users.
      rewriter.setInsertionPointAfter(op);
      // create new F16 constant op in that location
      ConstantOp new_const = rewriter.create<ConstantOp>(

    return {0, 2, 3, 1};
  } else {
    return {};
  }
}

// Shuffle elements in the `attr` according to the permutation. Optional
// `inner_size` allows to shuffle array attributes created from rank 2 tensors
// on outer dimension only.
ArrayAttr ShuffleArrayAttr(ArrayAttr attr, ArrayRef<int64_t> permutation,
                           int inner_size) {
  if (attr.empty()) return attr;

    builder->setInsertionPoint(terminator);
    auto yield = builder->create<tf_executor::YieldOp>(
        terminator->getLoc(), terminator->getOperands());
    terminator->erase();

    // Create new island for each region.
    builder->setInsertionPoint(island_op);
    auto execute_island = builder->create<tf_executor::IslandOp>(
        island_op.getLoc(), yield.getOperandTypes(),

      }
    }
  });
}

// Creates an instance of the default quant parameters pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateImportQuantStatsPass(
    OperationToName op_to_name, const std::string &stats_str) {
  auto pass = std::make_unique<ImportQuantStatsPass>(op_to_name);
  if (pass->ParseQuantStats(stats_str)) return nullptr;
  return pass;
}

constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

bool HasOutsideCompilationAttribute(Operation* op) {
  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))

  auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
  // Create an IslandOp that will wrap the IdentityOp. Add a control dependency
  // for the newly copied save function.
  auto wrapper_island_op = builder.create<IslandOp>(
      name_loc, TypeRange{main_file_prefix_arg.getType()},
      tf_executor::ControlType::get(&ctx), ValueRange(control_inputs));

  const std::string dump_file_path =
      tsl::io::JoinPath(*dump_dir, dump_file_name);
  TF_ASSIGN_OR_RETURN(std::unique_ptr<llvm::raw_ostream> file,
                      WritableFileWrapper::Create(dump_file_path));

  LOG(INFO) << "IR dump file created: " << dump_file_path;
  return file;
}

class PrinterConfig : public mlir::PassManager::IRPrinterConfig {
 public:
  explicit PrinterConfig(