if (block_arg.getParentRegion() != &func_op.getRegion()) {
        return false;
      }
      if (StringAttr attr = func_op.getArgAttrOfType<StringAttr>(
              block_arg.getArgNumber(), kTFDeviceAttr)) {
        return op_device == attr;
      }
    } else if (StringAttr attr = arg.getDefiningOp()->getAttrOfType<StringAttr>(
                   kDeviceAttr)) {
      return op_device == attr;
    }

  // pass the filter.
  MLIRContext context;
  auto id =
      StringAttr::get(&context, "//tensorflow/python/keras/keras_file.py");
  auto loc = FileLineColLoc::get(&context, id, 0, 0);
  auto id2 =
      StringAttr::get(&context, "//tensorflow/python/something/my_test.py");
  auto loc2 = FileLineColLoc::get(&context, id2, 0, 0);
  auto id3 = StringAttr::get(&context, "python/tensorflow/show_file.py");

    : public TypeStorage,
      public llvm::TrailingObjects<TFRTypeStorage, StringAttr> {
  using KeyTy = ArrayRef<StringAttr>;

  explicit TFRTypeStorage(unsigned num_attrs) : num_attrs(num_attrs) {}

  static TFRTypeStorage* construct(TypeStorageAllocator& allocator, KeyTy key) {
    // Allocate a new storage instance.
    auto byteSize = TFRTypeStorage::totalSizeToAlloc<StringAttr>(key.size());

      rewriter.getStringAttr(mhlo_custom_call.getCallTargetName()));

  if (auto bc = mhlo_custom_call.getBackendConfig()) {
    if (auto stringattr = mlir::dyn_cast_or_null<mlir::StringAttr>(*bc)) {
      tfl_custom.setCustomOptionAttr(
          TFL::ConstBytesAttr::get(rewriter.getContext(), stringattr));
    }
  } else {
    tfl_custom.setCustomOptionAttr(
        TFL::ConstBytesAttr::get(rewriter.getContext(), ""));
  }

    }
  }

  ChangeResult SetDevice(mlir::StringAttr device) {
    bool changed = (device != device_);
    device_ = device;
    return changed ? ChangeResult::Change : ChangeResult::NoChange;
  }

  bool hasDevice() const { return !!device_; }

  mlir::StringAttr getDevice() const { return device_; }

 private:
  mutable mlir::StringAttr device_;
};

class DeviceDataflowAnalysis

namespace mlir {
namespace TF {

using ::tensorflow::kValidDeviceTypes;

LogicalResult HasValidCompilationAndReplicationAttributes(Operation& op) {
  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);
  if (!replicate_attr && !compile_attr) return success();
  if (!replicate_attr || !compile_attr)

  for (auto op_attr : op_func_attr_array) {
    auto key_value = mlir::dyn_cast<mlir::ArrayAttr>(op_attr);
    if (!key_value || key_value.getValue().size() != 2 ||
        !mlir::isa<mlir::StringAttr>(key_value.getValue()[0]) ||
        !mlir::isa<mlir::StringAttr>(key_value.getValue()[1]))
      return op.emitOpError() << "each op_func_attr should be a key-value "
                                 "pair, where both the key and the value are "

  SmallVector<::mlir::NamedAttribute, 3> fields;

  auto padding_id = ::mlir::StringAttr::get(context, "padding");
  fields.emplace_back(padding_id, StringAttr::get(context, padding));

  auto pool_size_id = ::mlir::StringAttr::get(context, "pool_size");
  fields.emplace_back(pool_size_id, pool_size);

  auto strides_id = ::mlir::StringAttr::get(context, "strides");
  fields.emplace_back(strides_id, strides);

void ConstantOpDeviceAssignmentPass::runOnOperation() {
  ModuleOp module = getOperation();

  module.walk([&](TF::ConstOp op) {
    // Keep the ConstOp if the op already have the device attribute.
    if (StringAttr device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
      return WalkResult::advance();
    }
    OpBuilder builder(op);
    llvm::StringMap<mlir::Operation *> cloned_op_by_device;
    bool all_uses_replaced = true;

  // Parse node name from the `shared_name` attribute first. The variable v2 op
  // relies on the share name to look up from the TensorFlow's resource manager.
  StringAttr shared_name_attr = op->getAttrOfType<StringAttr>(kSharedNameAttr);
  if (shared_name_attr) {
    auto shared_name = StringRef(shared_name_attr.getValue());
    if (!shared_name.empty()) {
      return shared_name;
    }
  }