func.return
}

// -----

// Test case: Invalid multiple location information in a class attribute on VariableV2 op.

func.func @f() {
  // expected-error @+1 {{'tf.VariableV2' op expects only one named location in '_class' attribute, but got ["loc:@v1", "loc:@v2"]}}
  %val0 = "tf.VariableV2"() {_class = ["loc:@v1", "loc:@v2"], container = "", device = "", shape = #tf_type.shape<96>, shared_name = ""} : () -> tensor<96x!tf_type.f32ref>

// CHECK-LABEL: func @colocate_split_with_pred
func.func @colocate_split_with_pred() {
  // CHECK: Split
  // CHECK-SAME: _class = ["loc:@class"]
  tf_executor.graph {
    %c, %control0 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %a, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>)

    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Placeholder"
      }
    }
  }
}
node {
  name: "cond/MatMul/Switch_1"
  op: "Switch"
  input: "Const"
  input: "cond/pred_id"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Const"

    if (!shared_name.empty()) {
      return shared_name;
    }
  }
  // Attempt to parse "_class" attribute if there is no "shared_name"
  // attribute.
  ArrayAttr classes_attr = op->getAttrOfType<ArrayAttr>(kClassAttr);
  if (!classes_attr) {
    // Attempt to parse "_class" from the IdentityOp that follows VariableV2.
    // For read-only reference variables, IdentityOp should be the only user of

    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>
    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<100x50xf32>>>) -> tensor<100x50xf32>

    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "/job:worker/replica:0/task:1/device:CPU:0", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>
    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<100x50xf32>>>) -> tensor<100x50xf32>

module {
  // CHECK-LABEL: func @passesSingleIslandOp
  func.func @passesSingleIslandOp() {
    // CHECK: _class = ["loc:@class"]
    tf_executor.graph {
      %c, %control0 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
      %a, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>)

    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>
    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<100x50xf32>>>) -> tensor<100x50xf32>

}
node {
  name: "Variable/Assign"
  op: "Assign"
  input: "Variable"
  input: "Variable/initial_value"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Variable"
      }
    }
  }
  attr {
    key: "use_locking"
    value {
      b: true
    }
  }
  attr {
    key: "validate_shape"
    value {

func.func @remove_class_attribute(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) ->  tensor<f32> {
  // CHECK: "tf_device.cluster"
  // CHECK: "tf.Add"
  // CHECK-NOT: _class
  %1 = "tf_device.cluster"() ({
    %2 = "tf.Add"(%arg1, %arg2) {_class = ["loc:@while/BiasAdd_21/handle"]} : (tensor<f32>, tensor<f32>) -> tensor<f32>
    tf_device.return %2 : tensor<f32>