// devices, and empty StringRef will be returned. Control dependencies,
// NextIteration.Source -> NextIteration.Sink token dependencies, and
// LoopCond -> Switch data dependencies are ignored.
llvm::StringRef FindDeviceFromOperands(
    Operation& op,
    const llvm::DenseMap<Value, llvm::StringRef>& value_to_device) {
  llvm::StringRef new_device;

    mlir::ModuleOp module_op, llvm::StringRef device_type,
    xla::XlaComputation* xla_computation, bool use_tuple_args,
    bool enable_op_fallback, bool return_tuple,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

bool is_valid_special_tpu_op(
    std::vector<IslandOp>& ops, llvm::StringRef cluster_name,
    llvm::SmallDenseMap<llvm::StringRef, llvm::SmallDenseSet<Operation*>>&
        cluster_to_tpu_op_map) {
  for (IslandOp op : ops) {
    Operation* wrapped_op = &op.GetBody().front();
    std::optional<llvm::StringRef> wrapped_op_cluster_name =
        GetTpuClusterName(wrapped_op);

// Default platform for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
// Attribute required for running shape refinement pass enabled in XlaCallModule
// version 8 and above.
constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";

                              llvm::StringRef device_type,
                              bool enable_op_fallback,
                              llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
                                  custom_legalization_passes,
                              bool lower_to_xla_hlo,
                              llvm::StringRef module_name = llvm::StringRef()) {
  mlir::PassManager tf2xla(module_op.getContext());

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertMainFunctionPass)

  explicit InsertMainFunctionPass() = default;

  StringRef getArgument() const override {
    return "quant-insert-main-function";
  }

  StringRef getDescription() const override {
    return "Inserts the main function to the module.";
  }

  void runOnOperation() override;
};

  }
}

// Returns ID for identifying a resource.
std::tuple<llvm::StringRef, llvm::StringRef, llvm::StringRef> GetResourceKey(
    Operation* op) {
  llvm::StringRef device;
  if (auto attr = op->getAttrOfType<mlir::StringAttr>("device")) {
    device = attr.getValue();
  }

  llvm::StringRef container;
  if (auto attr = op->getAttrOfType<mlir::StringAttr>("container")) {

         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
             llvm::StringRef(
                 QuantTraitValues[QuantizationTrait::FullyQuantizable]);
}

// Returns the composite function name.
std::optional<StringRef> GetCompsiteFunctionName(Operation *op) {
  if (!IsCallToQuantizableLiftedFunction(op)) return std::nullopt;

  ModuleOp module = op->getParentOfType<ModuleOp>();
  // The StringRef cast is necessary before cxx14.
  if (failed(SetOnceModuleAttribute(
          StringRef(kGroupSizeAttrName.data(), kGroupSizeAttrName.size()),
          group_size, op, module))) {
    return failure();
  }
  if (failed(SetOnceModuleAttribute(
          StringRef(kGroupKeyAttrName.data(), kGroupKeyAttrName.size()),
          group_key, op, module))) {

// A TPU device for execution alongside its associated host CPU device.
struct TPUDeviceAndHost {
  TPUDeviceAndHost() = default;
  TPUDeviceAndHost(llvm::StringRef device, llvm::StringRef host)
      : device(device), host(host) {}

  std::string device;
  std::string host;
};

// TPU devices to be used for execution (e.g. devices for TPUExecute ops) and