auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kNumCoresPerReplicaAttr,
                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

  EXPECT_FALSE(HasModelParallelism(cluster));
}

  added_call_op->setAttr(kInferenceType, inference_type);
  added_func_op->setAttr(kDevice, device);
  added_func_op->setAttr(kInferenceType, inference_type);

  std::string function_name = absl::StrCat(interface_name.getValue().str(), "_",
                                           device.getValue().str(), "_",
                                           inference_type.getValue().str());

  // function name can be changed by those passes.
  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                   FlatSymbolRefAttr::get(builder.getStringAttr(func_name)));
  // Store the custom attribute to restore the function name when loading it
  // back in the post calibration stage. As mentioned above, the above entry
  // function attribute is not reliable.
  call_op->setAttr(kOriginalStablehloEntryFunctionAttrName,

        return func.emitError()
               << "failed to rename StableHLO function " << func.getSymName();
      }
    }
    if (is_main_func) {
      main_func_name = func.getSymNameAttr();
    }
    func->setAttr(kFromXlaCallModuleAttrName, builder.getUnitAttr());
  }
  if (!main_func_name) {
    return stablehlo_module.emitError()
           << "StableHLO module does not have an entry function";
  }
  return main_func_name;

  if (!interface_name.has_value()) {
    func.emitError("the func op does not have interface_name");
    return nullptr;
  }

  cloned_func->setAttr(
      kDevice, builder->getStringAttr(target_device_inference_type.hardware));
  cloned_func->setAttr(kInferenceType,
                       builder->getStringAttr(GetInferenceString(
                           target_device_inference_type.inference_type)));

      auto new_key = StringAttr::get(ctx, old_key.str() + "_dtoh_0");
      op->setAttr("key", new_key);
    } else if (isa<TF::XlaRecvFromHostOp>(op)) {
      StringRef old_key = op->getAttrOfType<StringAttr>("key").getValue();
      auto new_key = StringAttr::get(ctx, old_key.str() + "_htod_0");
      op->setAttr("key", new_key);
    }
  });
  return success();
}

  ConstOp zero = builder.create<ConstOp>(loc, zero_attr);
  zero->setAttr(kICIWeightDistributionMlirBridgeMarker,
                builder.getBoolAttr(true));
  ConstOp shape = builder.create<ConstOp>(loc, shape_attr);
  shape->setAttr(kICIWeightDistributionMlirBridgeMarker,
                 builder.getBoolAttr(true));
  FillOp fill = builder.create<FillOp>(loc, shape, zero);
  fill->setAttr(kICIWeightDistributionMlirBridgeMarker,

  // Both the device attribute and compile_device_type attribute should be
  // 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;

    if (auto device = res->getAttrOfType<StringAttr>("device")) {
      if (!device.getValue().empty())
        result_id->setAttr("device", device);
      else
        result_id->setAttr("device", compile_device_op);
    } else if (compile_device_op) {
      result_id->setAttr("device", compile_device_op);
    }
    res.getOutput().replaceAllUsesWith(compile_op->getResult(0));
  }

  new_call_op->setAttr(kEntryFuncAttrName,
                       rewriter.getStringAttr(new_ref_func_name.getValue()));
  new_call_op->setAttrs(call_op->getAttrs());
  new_call_op->removeAttr(rewriter.getStringAttr(kQuantTraitAttrName));

  FlatSymbolRefAttr new_func_name_attr =
      FlatSymbolRefAttr::get(rewriter.getContext(), new_ref_func_name);
  new_call_op->setAttr(kEntryFuncAttrName, new_func_name_attr);