for (auto fn : module->getOps<mlir::func::FuncOp>()) {
    for (auto& bb : fn) {
      for (auto& inst : bb) {
        auto attr_id = mlir::StringAttr::get(context, "value");
        if (auto attr = inst.getAttrOfType<mlir::ElementsAttr>(attr_id)) {
          mlir::Attribute rand_val;
          mlir::Type element_type = attr.getShapedType().getElementType();
          if (mlir::isa<mlir::IntegerType>(element_type)) {

      device = func.getArgAttrOfType<mlir::StringAttr>(arg.getArgNumber(),
                                                       kFuncDeviceAttr);
    } else {
      device = alias.getDefiningOp()->getAttrOfType<StringAttr>(kDeviceAttr);
    }

    if (!device) return false;
    tensorflow::DeviceNameUtils::ParsedName parsed_device;
    if (!tensorflow::DeviceNameUtils::ParseFullName(device.getValue().str(),

  return device_attr ? device_attr.getValue() : "";
}

// Extracts and canonicalizes the device attribute.
inline StringRef GetDeviceAttr(Operation* op) {
  auto device_attr = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  return device_attr ? device_attr.getValue() : "";
}

// Print operation with debug info (to get line number info for debugging)
void dump(StringRef message, Operation* op) {

// matches `initializer_type`.
bool IsInitializerType(func::FuncOp init_func_op, StringRef initializer_type) {
  auto init_type =
      init_func_op->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr);
  return init_type && init_type == initializer_type;
}

// Returns the initializer function whose tf_saved_model.initializer_type

LogicalResult SetOnceModuleAttribute(StringRef attr_name,
                                     IntegerAttr attr_value, Operation* op,
                                     ModuleOp& module) {
  const auto ex_attr_value = module->getAttrOfType<IntegerAttr>(attr_name);
  if (ex_attr_value == nullptr) {
    module->setAttr(attr_name, attr_value);
    return success();
  }
  if (ex_attr_value == attr_value) {
    return success();
  }

    return success();
  }
};

const char kDeviceAttr[] = "device";
const char kDeviceGpu[] = "GPU";

std::optional<std::string> GetDevice(mlir::Operation *op) {
  mlir::StringAttr device = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  if (!device || device.getValue().empty()) {
    return std::nullopt;
  }
  const std::string device_name = device.str();
  tensorflow::DeviceNameUtils::ParsedName parsed_name;

func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op,
                            const SymbolTable symbol_table) {
  const auto entry_function_symbol_ref =
      xla_call_module_op->getAttrOfType<FlatSymbolRefAttr>(kEntryFuncAttrName);

  return dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));
}

      llvm::dyn_cast<tf_device::LaunchOp>(preprocess_op->getParentOp());
  // Device of original launch looked up before moving the preprocessing ops
  // around.
  StringAttr device =
      original_launch_op->getAttrOfType<StringAttr>(kDeviceAttr);

  if (!device) {
    return original_launch_op->emitOpError()
           << "Launch op has an invalid device attribute.";
  }

  // Find the TPUCompile successor.

           << "]";

  return success();
}

bool VariableIsInitialized(TF::VarHandleOp var_handle_op) {
  auto is_variable_initialized =
      var_handle_op->getAttrOfType<BoolAttr>("_is_initialized");
  // Assume variable is initialized if attribute is not set.
  // There are paths that doesn't mark the variables. All variables
  // that doesn't have the attribute will be promoted.

    if (!legacy_graph_export) {
      std::string group_annotation = absl::StrCat(
          "r", std::to_string(replica_group_idx), ":", std::to_string(i));
      if (auto parallel_group_attr = replicate_op->getAttrOfType<StringAttr>(
              TF::kParallelExecAnnotation)) {
        // Extend the existing attribute so that nested parallel execution
        // structure is supported.