// CHECK-LABEL: func @device_test
func.func @device_test(%arg0: tensor<3x1xf32>) -> (tensor<3x3xf32>) attributes {tf.device = "xpu"} {

  // CHECK: device = "xpu"
  %0 = "tf.Const"() {value = dense<[[1.0, 2.0, 3.0]]> : tensor<1x3xf32>} : () -> tensor<1x3xf32>
  // CHECK: device = "xpu"

}

std::function<void(PassManager*)> legalizeTFPasses() {
  return [](PassManager* pm) {
    pm->addPass(mlir::mhlo::createLegalizeTFPass(
        /* legalize_chlo=*/true, llvm::StringRef("gpu/xpu"),
        /* prefer_tf2xla=*/false));
  };
}

TEST(XlaLegalizeTest, IllegalOp) {
  constexpr char kMlirIllegalOpStr[] = R"(

  func.return
}

// -----

func.func @invalid_device_type() {
  // expected-error@+1 {{'tf.opA' op has invalid '_xla_compile_device_type' value 'XPU'}}
  %0 = "tf.opA"() {_xla_compile_device_type = "XPU", _replication_info = "replicate", device = "/device:TPU:0", name = "name", is_stateless = true} : () -> tensor<i1>

  func.func @empty_compilation_attribute() {
    // expected-error@+1 {{'tf_device.cluster_func' op has invalid '_xla_compile_device_type' value 'XPU'}}