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(), ""));
  }

                         Operation* parent) {
  int num_replicas = metadata.getNumReplicas();
  int num_cores_per_replica = metadata.getNumCoresPerReplica();
  StringAttr tpu_replicate_attr =
      metadata->getAttrOfType<StringAttr>(TF::kTpuReplicateAttr);
  if (auto repinput = dyn_cast<TF::TPUReplicatedInputOp>(op)) {
    if (!ValidateReplicatedInput(repinput, num_replicas, tpu_replicate_attr))
      return false;
  }

  func_attrs.push_back(
      {StringAttr::get(context, "inputs"),
       StringAttr::get(context, absl::StrJoin(input_names, ","))});
  func_attrs.push_back(
      {StringAttr::get(context, "outputs"),
       StringAttr::get(context, absl::StrJoin(output_names, ","))});
  auto dictAttr = DictionaryAttr::get(context, func_attrs);
  main_func->setAttr(StringAttr::get(context, kEntryFunctionAttr), dictAttr);

    }
  }

  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)

// propagated through.
bool IsSupportedExecutorOp(Operation& op) {
  auto ops_have_same_device = [](Operation* lhs, Operation* rhs) {
    auto lhs_device_attr = lhs->getAttrOfType<StringAttr>(kDeviceAttr);
    auto rhs_device_attr = rhs->getAttrOfType<StringAttr>(kDeviceAttr);
    return (!lhs_device_attr && !rhs_device_attr) ||
           (lhs_device_attr && rhs_device_attr &&
            lhs_device_attr.getValue() == rhs_device_attr.getValue());

  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 "