check if the model require grad by output def. fix doc str. fix type …

…hint

deepmd_pt/model/model/atomic_model.py

@@ -5,6 +5,7 @@
 from typing import (
   Optional,
   Dict,
+  List,
 )
 from deepmd_utils.model_format import FittingOutputDef
 from deepmd_pt.model.task import Fitting
@@ -25,7 +26,7 @@ def get_rcut(self)->float:
     raise NotImplementedError
 
   @abstractmethod
-  def get_sel(self)->int:
+  def get_sel(self)->List[int]:
     raise NotImplementedError
 
   @abstractmethod
@@ -43,4 +44,27 @@ def forward_atomic(
   ) -> Dict[str, torch.Tensor]:
     raise NotImplementedError
 
+  def do_grad(
+      self, 
+      var_name: Optional[str] = None,
+  )->bool:
+    """Tell if the output variable `var_name` is differentiable.
+    if var_name is None, returns if any of the variable is differentiable.
     
+    """
+    odef = self.get_fitting_output_def()
+    if var_name is None:
+      require: List[bool] = []
+      for vv in odef.keys():
+        require.append(self.do_grad_(vv))
+      return any(require)
+    else:
+      return self.do_grad_(var_name)
+
+  def do_grad_(
+      self,
+      var_name: str,
+  )->bool:
+    """Tell if the output variable `var_name` is differentiable."""
+    assert var_name is not None
+    return self.get_fitting_output_def()[var_name].differentiable

deepmd_pt/model/model/dp_atomic_model.py

@@ -52,12 +52,6 @@ def __init__(
       sampled=None,
       **kwargs,
   ):
-    """Based on components, construct a DPA-1 model for energy.
-
-    Args:
-    - model_params: The Dict-like configuration with model options.
-    - sampled: The sampled dataset for stat.
-    """
     super().__init__()
     # Descriptor + Type Embedding Net (Optional)
     ntypes = len(type_map)
@@ -113,27 +107,32 @@ def __init__(
             self.coord_denoise_net = DenoiseNet(self.descriptor.dim_out, self.ntypes - 1, self.descriptor.dim_emb)
 
 
-  def get_fitting_net(self):
+  def get_fitting_net(self)->Fitting:
+    """Get the fitting net."""
     return (
       self.fitting_net
       if self.fitting_net is not None
       else self.coord_denoise_net
     )
 
   def get_fitting_output_def(self)->FittingOutputDef:
+    """Get the output def of the fitting net."""
     return (
       self.fitting_net.output_def() 
       if self.fitting_net is not None 
       else self.coord_denoise_net.output_def()
     )
 
-  def get_rcut(self):
+  def get_rcut(self)->float:
+    """Get the cut-off radius."""
     return self.rcut
 
-  def get_sel(self):
+  def get_sel(self)->List[int]:
+    """Get the neighbor selection."""
     return self.sel
 
-  def distinguish_types(self):
+  def distinguish_types(self)->bool:
+    """If distinguish different types by sorting."""
     return self.type_split
 
 
@@ -168,7 +167,7 @@ def forward_atomic(
     """
     nframes, nloc, nnei = nlist.shape
     atype = extended_atype[:, :nloc]
-    if self.grad_force:
+    if self.do_grad():
       extended_coord.requires_grad_(True)
     descriptor, env_mat, diff, rot_mat, sw = \
       self.descriptor(

deepmd_pt/model/model/ener.py

@@ -43,7 +43,7 @@ def forward(
             model_predict = {}
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
-            if self.grad_force:
+            if self.do_grad("energy"):
                 model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
                 model_predict["atomic_virial"] = model_ret["energy_derv_c"].squeeze(-3)
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-3)
@@ -69,7 +69,7 @@ def forward_lower(
             model_predict = {}
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
-            if self.grad_force:
+            if self.do_grad("energy"):
                 model_predict['extended_force'] = model_ret['energy_derv_r'].squeeze(-2)
                 model_predict['extended_virial'] = model_ret['energy_derv_c'].squeeze(-3)
             else:

deepmd_pt/model/model/make_model.py

@@ -31,6 +31,7 @@ def get_model_output_def(self):
                 self.get_fitting_output_def() 
             )
 
+        # cannot use the name forward. torch script does not work
         def forward_common(
             self,
             coord,

deepmd_pt/model/model/model.py

@@ -13,7 +13,7 @@ def __init__(self):
         """
         super(BaseModel, self).__init__()
 
-    def forward(self, **kwargs):
+    def forward(self, *args, **kwargs):
         """Model output.
         """
         raise NotImplementedError