// The sizes of indices, shared_name and container must be equal.
  int32_t indices_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("indices").size();
  int32_t shared_name_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("shared_name").size();
  int32_t container_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("container").size();

  if (!(indices_size == shared_name_size &&

  // propagated to the reduce function call.
  reduce_call->setAttr(kDeviceAttr,
                       reduce_dataset->getAttrOfType<StringAttr>(kDeviceAttr));
  reduce_call->setAttr(
      TF::kCompileDeviceTypeAttr,
      reduce_dataset->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr));

  SmallVector<Value, 4> if_returns;

  builder.create<TF::YieldOp>(loc,

ResourceHandle GetResourceHandle(Operation *op) {
  llvm::StringRef device;
  if (auto attr = op->getAttrOfType<StringAttr>("device")) {
    device = attr.getValue();
  }

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

  llvm::StringRef shared_name;
  if (auto attr = op->getAttrOfType<StringAttr>("shared_name")) {
    shared_name = attr.getValue();
  }

  // Duplicate of the logic in MultiDeviceFunctionBodyPlacer::BodyNodeDevice
  // LINT.IfChange
  auto device_id = StringAttr::get(call->getContext(), "device");
  auto caller_device = call->getAttrOfType<StringAttr>(device_id);
  if (!caller_device) return;

  DeviceNameUtils::ParsedName caller_parsed_device;
  if (!DeviceNameUtils::ParseFullName(caller_device.getValue().str(),

}

mlir::LogicalResult ExtractTfVersions(mlir::ModuleOp module,
                                      VersionDef* versions) {
  versions->Clear();
  auto version_attr =
      module->getAttrOfType<mlir::DictionaryAttr>("tf.versions");
  if (!version_attr) return mlir::failure();

  auto producer =
      mlir::dyn_cast_or_null<mlir::IntegerAttr>(version_attr.get("producer"));
  if (!producer) return mlir::failure();

  absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {
    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
    if (!cost_attr) return;
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (!op_key_attr) return;
    op_cost_map[op_key_attr.getInt()] = cost_attr.getInt();
  });
  return op_cost_map;
}

inline bool TFDataFormatIsNHWC(Operation *op) {
  auto attr = op->getAttrOfType<StringAttr>("data_format");
  return !attr || attr.getValue() == "NHWC";
}

// Returns true if the given TensorFlow op does not have a `data_format`
// attribute (then default to "NDHWC"), or its `data_format` attribute is
// "NDHWC".
inline bool TFDataFormatIsNDHWC(Operation *op) {
  auto attr = op->getAttrOfType<StringAttr>("data_format");

                  builder.getStringAttr(tensorflow::kTpuDevice));
    }

    // Convert `_XlaMustCompile`.
    if (op->hasAttr(tensorflow::kMustCompileAttr)) {
      bool must_compile_attr_val =
          op->getAttrOfType<BoolAttr>(tensorflow::kMustCompileAttr).getValue();
      op->removeAttr(tensorflow::kMustCompileAttr);
      if (!must_compile_attr_val) {
        if (op->hasAttr(tensorflow::kCompileDeviceTypeAttr)) {

  if (!call_op->hasAttr(kOriginalStablehloEntryFunctionAttrName)) {
    return;
  }

  const auto original_function_name = call_op->getAttrOfType<StringAttr>(
      kOriginalStablehloEntryFunctionAttrName);
  const auto current_function_name = call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);

  if (!original_function_name || !current_function_name) {
    return;
  }

  ModuleOp module = getOperation();

  module.walk([&](TF::ConstOp op) {
    // Keep the ConstOp if the op already have the device attribute.
    if (StringAttr device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
      return WalkResult::advance();
    }
    OpBuilder builder(op);
    llvm::StringMap<mlir::Operation *> cloned_op_by_device;
    bool all_uses_replaced = true;