mlir::cast<NameLoc>(locations.front()).getName().strref();
    StringRef op_type =
        op_type_with_suffix.substr(0, op_type_with_suffix.size() - 1);
    new_unit.set_op_type(op_type.str());

    if (isa<NameLoc>(locations.back())) {
      StringRef name_loc_id =
          mlir::cast<NameLoc>(locations.back()).getName().strref();
      set_node_and_func_name(new_unit, name_loc_id);
    } else {

  auto callsite_loc = llvm::dyn_cast<CallSiteLoc>(attr);

  if (!mlir::isa<NameLoc>(callsite_loc.getCaller())) return false;
  StringRef caller_name =
      mlir::cast<NameLoc>(callsite_loc.getCaller()).getName().strref();
  return caller_name.starts_with(kQuantizationUnitPrefix) &&
         caller_name.ends_with(kQuantizationUnitSuffix);
}

std::optional<QuantizationUnitLoc::QuantizationUnit>
FindQuantizationUnitFromLoc(Location loc) {

  absl::c_iota(spatial_dims, 1);
  bool has_dynamic_spatial_dim = absl::c_any_of(
      spatial_dims,
      [&input_shape](int64_t dim) { return input_shape.isDynamicDim(dim); });
  if (conv_padding.strref() == "SAME" && has_dynamic_spatial_dim) {
    return PadForDynamicShapedInputSamePadding(
        builder, loc, input, filter, input_zp_value, strides, dilations,
        conv_padding, padding, num_dims);
  }

    if (!locations.empty()) {
      // Skip locations for propagating op_type metadata.
      if (auto name_loc = mlir::dyn_cast<mlir::NameLoc>(locations[0])) {
        if (name_loc.getName().strref().ends_with(":")) {
          if (locations.size() == 2)
            return locations[1];
          else if (locations.size() > 2)
            return mlir::FusedLoc::get(
                fused_loc.getContext(),

    return attr.getName().strref().front() == '_';
  });
}

// Copies outside compilation attribute from `from` to `to`.
inline void CopyXlaOutsideCompilationAttributes(Operation *from,
                                                Operation *to) {
  CopyAttributes(from, to, [](const NamedAttribute &attr) {
    return attr.getName().strref() == kXlaOutsideCompilationAttr;
  });
}

          size_t num_locs = locations.size();
          // Skip the first location if it stores information for propagating
          // op_type metadata.
          if (num_locs > 0) {
            if (auto name_loc = mlir::dyn_cast<mlir::NameLoc>(locations[0])) {
              if (name_loc.getName().strref().ends_with(":")) {
                if (num_locs == 2) {

    return true;
  }
  std::string shard_string;
  if (op->hasAttr(kXLAShardingAttr)) {
    shard_string =
        op->getAttrOfType<StringAttr>(kXLAShardingAttr).strref().str();
  } else {
    shard_string = op->getAttrOfType<StringAttr>(kShardingAttr).strref().str();
  }
  xla::OpSharding sharding;
  if (!shard_string.empty() && sharding.ParseFromString(shard_string)) {
    // Only checking op with TUPLE sharding

    if (auto name = loc.dyn_cast<NameLoc>()) {
      return name.getName().strref();
    } else if (auto fused_name = loc.dyn_cast<FusedLoc>()) {
      for (auto sub_loc : fused_name.getLocations()) {
        if (auto named_sub_loc = sub_loc.dyn_cast<NameLoc>()) {
          return named_sub_loc.getName().strref();
        }
      }
    }
    return llvm::StringRef("");
  };

  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;
  mlir::dyn_cast_or_null<StringAttr>(entry_function_attr.get("inputs"))
      .strref()
      .split(inputs, /*Separator=*/",");

  return inputs;
}

void ConvertMainArgAttrs(func::FuncOp main_func_op, const int arg_idx,
                         const StringRef index_path) {

    if (!DeviceNameUtils::ParseFullName(name.str(), &device))
      return op->emitOpError(
          llvm::formatv("bad '{0}' attribute, '{1}', not a valid device",
                        kDevicesAttr, name.strref()));

    if (auto gpu_metadata = mlir::dyn_cast<mlir::TF::GpuDeviceMetadata>(attr)) {
      devices->AddGpuDevice(device, gpu_metadata);
    } else {
      devices->AddDevice(device);
    }
  }