fix type annotations and refine docstrings

Author: HydrogenSulfate

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -134,8 +134,16 @@ def get_type_map(self) -> list[str]:
         """Get the type map."""
         return self.type_map
 
-    def get_buffer_type_map(self) -> list[str]:
-        """Get the type map."""
+    def get_buffer_type_map(self) -> paddle.Tensor:
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_compute_stats_distinguish_types(self) -> bool:
@@ -602,7 +610,7 @@ def _store_out_stat(
     def get_ntypes(self):
         return len(self.type_map)
 
-    def get_buffer_ntypes(self):
+    def get_buffer_ntypes(self) -> paddle.Tensor:
         return self.buffer_ntypes
 
     def _fetch_out_stat(

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -160,15 +160,23 @@ def get_sel(self) -> list[int]:
         return self.sel
 
     def get_buffer_type_map(self) -> paddle.Tensor:
-        """Get the neighbor selection."""
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Get the cut-off radius."""
+        """Get the cut-off radius as a buffer-style Tensor."""
         return self.descriptor.get_buffer_rcut()
 
     def get_buffer_sel(self) -> paddle.Tensor:
-        """Get the neighbor selection."""
+        """Get the neighbor selection as a buffer-style Tensor."""
         return self.descriptor.get_buffer_sel()
 
     def set_case_embd(self, case_idx: int) -> None:
@@ -400,7 +408,7 @@ def get_dim_fparam(self) -> int:
         return self.fitting_net.get_dim_fparam()
 
     def get_buffer_dim_fparam(self) -> paddle.Tensor:
-        """Get the number (dimension) of frame parameters of this atomic model."""
+        """Get the number (dimension) of frame parameters of this atomic model as a buffer-style Tensor."""
         return self.fitting_net.get_buffer_dim_fparam()
 
     def has_default_fparam(self) -> bool:
@@ -412,7 +420,7 @@ def get_dim_aparam(self) -> int:
         return self.fitting_net.get_dim_aparam()
 
     def get_buffer_dim_aparam(self) -> paddle.Tensor:
-        """Get the number (dimension) of atomic parameters of this atomic model."""
+        """Get the number (dimension) of atomic parameters of this atomic model as a buffer-style Tensor."""
         return self.fitting_net.get_buffer_dim_aparam()
 
     def get_sel_type(self) -> list[int]:

deepmd/pd/model/descriptor/dpa1.py

@@ -325,16 +325,16 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.se_atten.get_rcut()
 
-    def get_buffer_rcut(self) -> float:
-        """Returns the cut-off radius."""
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.se_atten.get_buffer_rcut()
 
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.se_atten.get_rcut_smth()
 
-    def get_buffer_rcut_smth(self) -> float:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+    def get_buffer_rcut_smth(self) -> paddle.Tensor:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.se_atten.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:
@@ -354,7 +354,15 @@ def get_type_map(self) -> list[str]:
         return self.type_map
 
     def get_buffer_type_map(self) -> paddle.Tensor:
-        """Get the name to each type of atoms."""
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:

deepmd/pd/model/descriptor/dpa2.py

@@ -345,7 +345,7 @@ def get_buffer_rcut(self) -> paddle.Tensor:
         return self.repinit.get_buffer_rcut()
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.repinit.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:

deepmd/pd/model/descriptor/dpa3.py

@@ -230,11 +230,11 @@ def get_rcut_smth(self) -> float:
         return self.rcut_smth
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.repflows.get_buffer_rcut()
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.repflows.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:
@@ -246,7 +246,7 @@ def get_sel(self) -> list[int]:
         return self.sel
 
     def get_buffer_sel(self) -> paddle.Tensor:
-        """Returns the number of selected atoms for each type."""
+        """Returns the number of selected atoms for each type as a buffer-style Tensor."""
         return self.repflows.get_sel()
 
     def get_ntypes(self) -> int:
@@ -258,7 +258,15 @@ def get_type_map(self) -> list[str]:
         return self.type_map
 
     def get_buffer_type_map(self) -> paddle.Tensor:
-        """Get the name to each type of atoms."""
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:

deepmd/pd/model/descriptor/repflows.py

@@ -361,7 +361,7 @@ def get_rcut(self) -> float:
         return self.e_rcut
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.buffer_rcut
 
     def get_rcut_smth(self) -> float:
@@ -377,7 +377,7 @@ def get_sel(self) -> list[int]:
         return self.sel
 
     def get_buffer_sel(self) -> paddle.Tensor:
-        """Returns the number of selected atoms for each type."""
+        """Returns the number of selected atoms for each type as a buffer-style Tensor."""
         return self.buffer_sel
 
     def get_ntypes(self) -> int:

deepmd/pd/model/descriptor/repformers.py

@@ -330,11 +330,11 @@ def get_rcut_smth(self) -> float:
         return self.rcut_smth
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.buffer_rcut
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int:

deepmd/pd/model/descriptor/se_a.py

@@ -128,11 +128,11 @@ def get_rcut_smth(self) -> float:
         return self.sea.get_rcut_smth()
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.sea.get_buffer_rcut()
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.sea.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:
@@ -152,7 +152,15 @@ def get_type_map(self) -> list[str]:
         return self.type_map
 
     def get_buffer_type_map(self) -> paddle.Tensor:
-        """Get the name to each type of atoms."""
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:
@@ -536,11 +544,11 @@ def get_rcut_smth(self) -> float:
         return self.rcut_smth
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.buffer_rcut
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int:

deepmd/pd/model/descriptor/se_atten.py

@@ -285,11 +285,11 @@ def get_rcut_smth(self) -> float:
         return self.rcut_smth
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.buffer_rcut
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int:

deepmd/pd/model/descriptor/se_t_tebd.py

@@ -214,7 +214,15 @@ def get_type_map(self) -> list[str]:
         return self.type_map
 
     def get_buffer_type_map(self) -> paddle.Tensor:
-        """Get the name to each type of atoms."""
+        """
+        Return the type map as a buffer-style Tensor for JIT saving.
+
+        The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string
+        (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`,
+        and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int.
+
+        This format allows the type map to be serialized as a raw byte buffer during JIT model saving.
+        """
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:
@@ -631,11 +639,11 @@ def get_rcut_smth(self) -> float:
         return self.rcut_smth
 
     def get_buffer_rcut(self) -> paddle.Tensor:
-        """Returns the cut-off radius."""
+        """Returns the cut-off radius as a buffer-style Tensor."""
         return self.buffer_rcut
 
     def get_buffer_rcut_smth(self) -> paddle.Tensor:
-        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int: