diff --git a/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc b/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc
index 272fab9cd1c..bce0ed4a33d 100644
--- a/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc
+++ b/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc
@@ -55,8 +55,10 @@ llvm::StringRef OpOrArgNameMapper::GetUniqueName(llvm::StringRef prefix) {
   // to be unique.
   auto& val = prefix_it.first->second;
   llvm::SmallString<64> probe_name(prefix);
+  probe_name.append(GetSuffixSeparator());
+  const int probe_prefix_size = probe_name.size();
   while (true) {
-    probe_name.resize(prefix.size());
+    probe_name.resize(probe_prefix_size);
     // TODO(jpienaar): Subtract one so that the initial suffix is 0 instead
     // of 1.
     // TODO(jpienaar): Switch to radix 36 and update tests.
diff --git a/tensorflow/compiler/mlir/op_or_arg_name_mapper.h b/tensorflow/compiler/mlir/op_or_arg_name_mapper.h
index 108496e2283..6a52d13fbc0 100644
--- a/tensorflow/compiler/mlir/op_or_arg_name_mapper.h
+++ b/tensorflow/compiler/mlir/op_or_arg_name_mapper.h
@@ -64,6 +64,9 @@ class OpOrArgNameMapper {
     return op_or_val_to_name_;
   }
 
+  // Returns the separator used before uniqueing suffix.
+  virtual llvm::StringRef GetSuffixSeparator() { return ""; }
+
  private:
   // Returns name from the location of the operation or value.
   virtual std::string GetName(OpOrVal op_or_val) = 0;