if (HasDevice(split) || split->getAttrOfType<ArrayAttr>(kClassAttr))
        return WalkResult::advance();
      for (Operation* pred : IslandPredecessors(split)) {
        if (auto colocation_classes =
                pred->getAttrOfType<ArrayAttr>(kClassAttr)) {
          split->setAttr(kClassAttr, colocation_classes);
          return WalkResult::advance();
        }
      }
    }

#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

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

  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
    // VariableV2.
    auto identity_op = op->getUsers().begin();
    classes_attr = identity_op->getAttrOfType<ArrayAttr>(kClassAttr);
    if (!classes_attr) {