namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kClassAttr[] = "_class";

#define GEN_PASS_DEF_TPUCLEANUPCLUSTERATTRIBUTESPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

class TPUCleanupClusterAttributesPass
    : public impl::TPUCleanupClusterAttributesPassBase<
          TPUCleanupClusterAttributesPass> {
 public:

    %7:3 = tf_executor.Switch %3#0, %6#0 : (tensor<*xi32>, tensor<*xi1>) -> (tensor<*xi32>, tensor<*xi32>, !tf_executor.control) {device = "", T = "tfdtype$DT_INT32", _class = ["loc:@while/Merge"]} loc("while/Switch")
    %8:2 = tf_executor.Exit %7#1 : tensor<*xi32> {device = "", T = "tfdtype$DT_INT32"} loc("while/Exit")

        "tf_executor.yield"(%0) : (tensor<f32>) -> ()
      }) : () -> (tensor<f32>, !tf_executor.control)
      %outputs_16:2, %control_17 = "tf_executor._SwitchN"(%outputs_14, %outputs) {T = f32, _class = ["Case/input_0"], device = "", num_outs = 2 : i64} : (tensor<f32>, tensor<i32>) -> (tensor<*xf32>, tensor<*xf32>, !tf_executor.control)
      %outputs_18, %control_19 = "tf_executor.island"() ({

}
node {
  name: "Case/Case/input_0/_7"
  op: "_SwitchN"
  input: "Case/input_0"
  input: "Case/branch_index"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "Case/input_0"
      }
    }
  }
  attr {
    key: "num_outs"
    value {
      i: 2
    }
  }
}
node {
  name: "Case/pivot_0/_4"

  NameAttrList name;
  name.set_name(def.op());
  *name.mutable_attr() = def.attr();
  // Remove the "_class" attribute from the attribute set used to create the
  // compilation cache key. This attribute is information for the colocator
  // and causes false uniqueness between nodes.
  name.mutable_attr()->erase("_class");
  return CompileImpl(compile_options, options, name, args, CompileScope::kOp,

  EXPECT_EQ(constraints[1], "loc:@c1");

  Scope with_colocate = root.ColocateWith(c3).ColocateWith(c4);
  auto c6 = Const(with_colocate.WithOpName("c6").ClearColocation(), 7);
  EXPECT_FALSE(c6.op().node()->attrs().Find("_class"));
}

TEST(CCOpTest, TemplatedConst) {
  Scope root = Scope::NewRootScope();
  auto c1 = ops::Const<float>(root, {{3, 2}, {-1, 0}});
  TF_EXPECT_OK(root.status());

  Tensor out;

}
node {
  name: "MobilenetV1/Conv2d_0/weights/read"
  op: "Identity"
  input: "MobilenetV1/Conv2d_0/weights"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@MobilenetV1/Conv2d_0/weights"
      }
    }
  }
}
node {
  name: "MobilenetV1/MobilenetV1/Conv2d_0/Conv2D"
  op: "Conv2D"
  input: "input"

  %48 = "tf.Identity"(%47) {device = ""} : (tensor<i1>) -> tensor<i1>
  %49 = "tf.Identity"(%36#1) {_class = ["loc:@WhitespaceTokenize/WhitespaceTokenizeWithOffsets"], device = ""} : (tensor<?xi64>) -> tensor<?xi64>

    The colocation class (`_class`) of the `Split` is set to the same class as
    its predecessor:

    ```mlir
    %outputs1:2, %control1 = tf_executor.island wraps "tf.IteratorGetNext"(%arg)
      {_class = ["loc:@dataset_iterator_1"]}
    %outputs2:2, %control2 = tf_executor.island wraps "tf.Split"(%outputs0, %outputs1#1)
      {_class = ["loc:@dataset_iterator_1", num_split = 2 : i32}
    ```

        sourceWithFileSuffixForProject("java", "impl", "class F {}")
        sourceWithFileSuffixForProject("groovy", "impl", "class G {}")
        run ":impl:compileGroovy"

        when:
        aClass.text = "class A { void m1() {}; void m2() {}; }"
        eClass.text = "class E { boolean isCacheEnabled = D.cache; garbage }"