message: Address failures in compiler tests when oneDNN is enabled.
diff: @@ -1072,17 +1072,22 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
  // Substring that should be checked for in device name for CPU device.
  const char* const kCPUDeviceSubStr = "CPU";
  + const char* const kXLACPUDeviceSubStr = "XLA_CPU";
  - // If Op has been specifically assigned to a non-CPU device, then No.
  + // If Op has been specifically assigned to a non-CPU or XLA_CPU device, then
  + // No.
  if (!n->assigned_device_name().empty() &&
  - !absl::StrContains(n->assigned_device_name(), kCPUDeviceSubStr)) {
  + (!absl::StrContains(n->assigned_device_name(), kCPUDeviceSubStr) ||
  + absl::StrContains(n->assigned_device_name(), kXLACPUDeviceSubStr))) {
  result = false;
  reason = "Op has been assigned a runtime device that is not CPU.";
  }
  - // If user has specifically assigned this op to a non-CPU device, then No.
  + // If user has specifically assigned this op to a non-CPU or XLA_CPU device,
  + // then No.
  if (!n->def().device().empty() &&
  - !absl::StrContains(n->def().device(), kCPUDeviceSubStr)) {
  + (!absl::StrContains(n->def().device(), kCPUDeviceSubStr) ||
  + absl::StrContains(n->def().device(), kXLACPUDeviceSubStr))) {
  result = false;
  reason = "User has assigned a device that is not CPU.";
  }
  