feat: add has_message_passing API

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -80,6 +80,10 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return self.descriptor.has_message_passing()
+
     def forward_atomic(
         self,
         extended_coord: np.ndarray,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -92,6 +92,10 @@
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return any(model.has_message_passing() for model in self.models)
+
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -115,6 +115,10 @@ def mixed_types(self) -> bool:
             """
             pass
 
+        @abstractmethod
+        def has_message_passing(self) -> bool:
+            """Returns whether the descriptor has message passing."""
+
         @abstractmethod
         def fwd(
             self,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -131,6 +131,10 @@
         # to match DPA1 and DPA2.
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return False
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/dpmodel/descriptor/descriptor.py

@@ -125,3 +125,7 @@ def call(
     ):
         """Calculate DescriptorBlock."""
         pass
+
+    @abstractmethod
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""

deepmd/dpmodel/descriptor/dpa1.py

@@ -349,6 +349,10 @@
         """
         return self.se_atten.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.se_atten.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
@@ -885,6 +889,10 @@
             sw,
         )
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False
+
 
 class NeighborGatedAttention(NativeOP):
     def __init__(

deepmd/dpmodel/descriptor/dpa2.py

@@ -524,6 +524,12 @@
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(
+            [self.repinit.has_message_passing(), self.repformers.has_message_passing()]
+        )
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/hybrid.py

@@ -138,6 +138,10 @@
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(descrpt.has_message_passing() for descrpt in self.descrpt_list)
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix. All descriptors should be the same."""
         all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -95,6 +95,10 @@ def mixed_types(self) -> bool:
             """
             pass
 
+        @abstractmethod
+        def has_message_passing(self) -> bool:
+            """Returns whether the descriptor has message passing."""
+
         @abstractmethod
         def get_env_protection(self) -> float:
             """Returns the protection of building environment matrix."""

deepmd/dpmodel/descriptor/repformers.py

@@ -387,6 +387,10 @@
         rot_mat = np.transpose(h2g2, (0, 1, 3, 2))
         return g1, g2, h2, rot_mat.reshape(-1, nloc, self.dim_emb, 3), sw
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return True
+
 
 # translated by GPT and modified
 def get_residual(

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -247,6 +247,10 @@ def mixed_types(self):
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_r.py

@@ -205,6 +205,10 @@ def mixed_types(self):
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_t.py

@@ -185,6 +185,10 @@
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/model/make_model.py

@@ -471,6 +471,10 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        def has_message_passing(self) -> bool:
+            """Returns whether the model has message passing."""
+            return self.atomic_model.has_message_passing()
+
         def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()

deepmd/pt/entrypoints/main.py

@@ -303,10 +303,7 @@ def train(FLAGS):
 
 def freeze(FLAGS):
     model = torch.jit.script(inference.Tester(FLAGS.model, head=FLAGS.head).model)
-    if '"type": "dpa2"' in model.get_model_def_script():
-        extra_files = {"type": "dpa2"}
-    else:
-        extra_files = {"type": "else"}
+    extra_files = {}
     torch.jit.save(
         model,
         FLAGS.output,

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -95,6 +95,10 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return self.descriptor.has_message_passing()
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -104,6 +104,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return any(model.has_message_passing() for model in self.models)
+
     def get_out_bias(self) -> torch.Tensor:
         return self.out_bias
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -160,6 +160,10 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return False
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/descriptor/descriptor.py

@@ -169,6 +169,10 @@ def forward(
         """Calculate DescriptorBlock."""
         pass
 
+    @abstractmethod
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+
 
 def make_default_type_embedding(
     ntypes,

deepmd/pt/model/descriptor/dpa1.py

@@ -342,6 +342,10 @@ def mixed_types(self) -> bool:
         """
         return self.se_atten.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.se_atten.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()

deepmd/pt/model/descriptor/dpa2.py

@@ -294,6 +294,12 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(
+            [self.repinit.has_message_passing(), self.repformers.has_message_passing()]
+        )
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         # the env_protection of repinit is the same as that of the repformer

deepmd/pt/model/descriptor/hybrid.py

@@ -143,6 +143,10 @@ def mixed_types(self):
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(descrpt.has_message_passing() for descrpt in self.descrpt_list)
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix. All descriptors should be the same."""
         all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]

deepmd/pt/model/descriptor/repformers.py

@@ -548,3 +548,7 @@ def get_stats(self) -> Dict[str, StatItem]:
                 "The statistics of the descriptor has not been computed."
             )
         return self.stats
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return True

deepmd/pt/model/descriptor/se_a.py

@@ -147,6 +147,10 @@ def mixed_types(self):
         """
         return self.sea.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.sea.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.sea.get_env_protection()
@@ -674,3 +678,7 @@ def forward(
             None,
             sw,
         )
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False

deepmd/pt/model/descriptor/se_atten.py

@@ -581,6 +581,10 @@ def forward(
             sw,
         )
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False
+
 
 class NeighborGatedAttention(nn.Module):
     def __init__(

deepmd/pt/model/descriptor/se_r.py

@@ -170,6 +170,10 @@
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/pt/model/descriptor/se_t.py

@@ -174,6 +174,10 @@
         """
         return self.seat.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.seat.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.seat.get_env_protection()
@@ -687,3 +691,7 @@
             None,
             sw,
         )
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False

deepmd/pt/model/model/frozen.py

@@ -106,6 +106,11 @@
         """
         return self.model.mixed_types()
 
+    @torch.jit.export
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.model.has_message_passing()
+
     @torch.jit.export
     def forward(
         self,

deepmd/pt/model/model/make_model.py

@@ -531,6 +531,11 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        @torch.jit.export
+        def has_message_passing(self) -> bool:
+            """Returns whether the model has message passing."""
+            return self.atomic_model.has_message_passing()
+
         def forward(
             self,
             coord,

source/api_cc/src/DeepPotPT.cc

@@ -59,13 +59,7 @@ void DeepPotPT::init(const std::string& model,
   }
   std::unordered_map<std::string, std::string> metadata = {{"type", ""}};
   module = torch::jit::load(model, device, metadata);
-  // TODO: This should be fixed after implement api to decide whether need to
-  // message passing and rename this metadata
-  if (metadata["type"] == "dpa2") {
-    do_message_passing = 1;
-  } else {
-    do_message_passing = 0;
-  }
+  do_message_passing = module.run_method("has_message_passing").toBool();
   torch::jit::FusionStrategy strategy;
   strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}};
   torch::jit::setFusionStrategy(strategy);