diff --git a/xprof/utils/hlo_proto_map.h b/xprof/utils/hlo_proto_map.h
index 221c4cc45..ea614b1f3 100644
--- a/xprof/utils/hlo_proto_map.h
+++ b/xprof/utils/hlo_proto_map.h
@@ -49,10 +49,14 @@ class HloProtoMap {
   auto begin() const { return hlo_protos_by_program_id_.begin(); }
   auto end() const { return hlo_protos_by_program_id_.end(); }
 
-  bool contains(absl::string_view name) const {
+  bool ContainsOptimizedModule(absl::string_view name) const {
     return hlo_protos_by_name_.contains(name);
   }
 
+  bool ContainsOriginalModule(absl::string_view name) const {
+    return original_hlo_protos_by_name_.contains(name);
+  }
+
   bool contains(uint64_t program_id) const {
     return hlo_protos_by_program_id_.contains(program_id);
   }