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());

         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;

limitations under the License.
==============================================================================*/

#include <memory>
#include <optional>
#include <string>
#include <utility>

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project

    absl::flat_hash_set<Node*>* control_ret_nodes) {
  mlir::Block& block = function.front();

  // Extract input & output names if set.
  llvm::SmallVector<llvm::StringRef, 2> input_names;
  llvm::SmallVector<llvm::StringRef, 2> output_names;
  llvm::SmallVector<llvm::StringRef, 2> unique_output_names;
  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {

// `tf.backend_config` is a DictionaryAttr, JAX2TF sets the value of its
// i64 attribute `called_index` to the TF function's name.
constexpr llvm::StringRef kTfBackendConfigAttrName = "tf.backend_config";
constexpr llvm::StringRef kCalledIndexAttrName = "called_index";
constexpr llvm::StringRef kCalledFuncAttrName = "called_func";

// Deserialize the StableHLO module embedded in XlaCallModuleOp's module
// attribute.

#include "flatbuffers/vector.h"  // from @flatbuffers
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/ErrorHandling.h"
#include "mlir/IR/AttrTypeSubElements.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      MergeInitializerFunctionOpsToMainPass)

  explicit MergeInitializerFunctionOpsToMainPass() = default;

  StringRef getArgument() const override {
    return "quant-merge-initializer-function-ops-to-main";
  }

  StringRef getDescription() const override {
    return "Moves all ops from the initializer functions to the main function. "