remove paddle.jit.export

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -4,10 +4,7 @@
 import logging
 from typing import (
     Callable,
-    Dict,
-    List,
     Optional,
-    Tuple,
     Union,
 )
 
@@ -76,11 +73,11 @@ class BaseAtomicModel(paddle.nn.Layer, BaseAtomicModel_):
 
     def __init__(
         self,
-        type_map: List[str],
-        atom_exclude_types: List[int] = [],
-        pair_exclude_types: List[Tuple[int, int]] = [],
+        type_map: list[str],
+        atom_exclude_types: list[int] = [],
+        pair_exclude_types: list[tuple[int, int]] = [],
         rcond: Optional[float] = None,
-        preset_out_bias: Optional[Dict[str, paddle.Tensor]] = None,
+        preset_out_bias: Optional[dict[str, paddle.Tensor]] = None,
     ):
         paddle.nn.Layer.__init__(self)
         BaseAtomicModel_.__init__(self)
@@ -93,7 +90,7 @@ def __init__(
     def init_out_stat(self):
         """Initialize the output bias."""
         ntypes = self.get_ntypes()
-        self.bias_keys: List[str] = list(self.fitting_output_def().keys())
+        self.bias_keys: list[str] = list(self.fitting_output_def().keys())
         self.max_out_size = max(
             [self.atomic_output_def()[kk].size for kk in self.bias_keys]
         )
@@ -122,14 +119,13 @@ def __getitem__(self, key):
         else:
             raise KeyError(key)
 
-    # @paddle.jit.export
-    def get_type_map(self) -> List[str]:
+    def get_type_map(self) -> list[str]:
         """Get the type map."""
         return self.type_map
 
     def reinit_atom_exclude(
         self,
-        exclude_types: List[int] = [],
+        exclude_types: list[int] = [],
     ):
         self.atom_exclude_types = exclude_types
         if exclude_types == []:
@@ -139,7 +135,7 @@ def reinit_atom_exclude(
 
     def reinit_pair_exclude(
         self,
-        exclude_types: List[Tuple[int, int]] = [],
+        exclude_types: list[tuple[int, int]] = [],
     ):
         self.pair_exclude_types = exclude_types
         if exclude_types == []:
@@ -194,8 +190,8 @@ def forward_common_atomic(
         mapping: Optional[paddle.Tensor] = None,
         fparam: Optional[paddle.Tensor] = None,
         aparam: Optional[paddle.Tensor] = None,
-        comm_dict: Optional[Dict[str, paddle.Tensor]] = None,
-    ) -> Dict[str, paddle.Tensor]:
+        comm_dict: Optional[dict[str, paddle.Tensor]] = None,
+    ) -> dict[str, paddle.Tensor]:
         """Common interface for atomic inference.
 
         This method accept extended coordinates, extended atom typs, neighbor list,
@@ -277,8 +273,8 @@ def forward(
         mapping: Optional[paddle.Tensor] = None,
         fparam: Optional[paddle.Tensor] = None,
         aparam: Optional[paddle.Tensor] = None,
-        comm_dict: Optional[Dict[str, paddle.Tensor]] = None,
-    ) -> Dict[str, paddle.Tensor]:
+        comm_dict: Optional[dict[str, paddle.Tensor]] = None,
+    ) -> dict[str, paddle.Tensor]:
         return self.forward_common_atomic(
             extended_coord,
             extended_atype,
@@ -290,7 +286,7 @@ def forward(
         )
 
     def change_type_map(
-        self, type_map: List[str], model_with_new_type_stat=None
+        self, type_map: list[str], model_with_new_type_stat=None
     ) -> None:
         """Change the type related params to new ones, according to `type_map` and the original one in the model.
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
@@ -355,7 +351,7 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel":
 
     def compute_or_load_stat(
         self,
-        merged: Union[Callable[[], List[dict]], List[dict]],
+        merged: Union[Callable[[], list[dict]], list[dict]],
         stat_file_path: Optional[DPPath] = None,
     ):
         """
@@ -378,7 +374,7 @@ def compute_or_load_stat(
 
     def compute_or_load_out_stat(
         self,
-        merged: Union[Callable[[], List[dict]], List[dict]],
+        merged: Union[Callable[[], list[dict]], list[dict]],
         stat_file_path: Optional[DPPath] = None,
     ):
         """
@@ -405,7 +401,7 @@ def compute_or_load_out_stat(
 
     def apply_out_stat(
         self,
-        ret: Dict[str, paddle.Tensor],
+        ret: dict[str, paddle.Tensor],
         atype: paddle.Tensor,
     ):
         """Apply the stat to each atomic output.
@@ -521,7 +517,7 @@ def _default_std(self):
 
     def _varsize(
         self,
-        shape: List[int],
+        shape: list[int],
     ) -> int:
         output_size = 1
         len_shape = len(shape)
@@ -533,7 +529,7 @@ def _get_bias_index(
         self,
         kk: str,
     ) -> int:
-        res: List[int] = []
+        res: list[int] = []
         for i, e in enumerate(self.bias_keys):
             if e == kk:
                 res.append(i)
@@ -542,8 +538,8 @@ def _get_bias_index(
 
     def _store_out_stat(
         self,
-        out_bias: Dict[str, paddle.Tensor],
-        out_std: Dict[str, paddle.Tensor],
+        out_bias: dict[str, paddle.Tensor],
+        out_std: dict[str, paddle.Tensor],
         add: bool = False,
     ):
         ntypes = self.get_ntypes()
@@ -563,8 +559,8 @@ def _store_out_stat(
 
     def _fetch_out_stat(
         self,
-        keys: List[str],
-    ) -> Tuple[Dict[str, paddle.Tensor], Dict[str, paddle.Tensor]]:
+        keys: list[str],
+    ) -> tuple[dict[str, paddle.Tensor], dict[str, paddle.Tensor]]:
         ret_bias = {}
         ret_std = {}
         ntypes = self.get_ntypes()

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -3,8 +3,6 @@
 import functools
 import logging
 from typing import (
-    Dict,
-    List,
     Optional,
 )
 
@@ -52,7 +50,7 @@ def __init__(
         self,
         descriptor,
         fitting,
-        type_map: List[str],
+        type_map: list[str],
         **kwargs,
     ):
         super().__init__(type_map, **kwargs)
@@ -66,7 +64,7 @@ def __init__(
         super().init_out_stat()
 
         # register 'type_map' as buffer
-        def _string_to_array(s: str) -> List[int]:
+        def _string_to_array(s: str) -> list[int]:
             return [ord(c) for c in s]
 
         self.register_buffer(
@@ -107,7 +105,6 @@ def _string_to_array(s: str) -> List[int]:
         )
         self.buffer_aparam_nall.name = "buffer_aparam_nall"
 
-    # @paddle.jit.export
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return (
@@ -116,12 +113,11 @@ def fitting_output_def(self) -> FittingOutputDef:
             else self.coord_denoise_net.output_def()
         )
 
-    # @paddle.jit.export
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return self.rcut
 
-    def get_sel(self) -> List[int]:
+    def get_sel(self) -> list[int]:
         """Get the neighbor selection."""
         return self.sel
 
@@ -138,7 +134,7 @@ def mixed_types(self) -> bool:
         return self.descriptor.mixed_types()
 
     def change_type_map(
-        self, type_map: List[str], model_with_new_type_stat=None
+        self, type_map: list[str], model_with_new_type_stat=None
     ) -> None:
         """Change the type related params to new ones, according to `type_map` and the original one in the model.
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
@@ -199,8 +195,8 @@ def forward_atomic(
         mapping: Optional[paddle.Tensor] = None,
         fparam: Optional[paddle.Tensor] = None,
         aparam: Optional[paddle.Tensor] = None,
-        comm_dict: Optional[Dict[str, paddle.Tensor]] = None,
-    ) -> Dict[str, paddle.Tensor]:
+        comm_dict: Optional[dict[str, paddle.Tensor]] = None,
+    ) -> dict[str, paddle.Tensor]:
         """Return atomic prediction.
 
         Parameters
@@ -300,7 +296,7 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.fitting_net.get_dim_aparam()
 
-    def get_sel_type(self) -> List[int]:
+    def get_sel_type(self) -> list[int]:
         """Get the selected atom types of this model.
 
         Only atoms with selected atom types have atomic contribution

deepmd/pd/model/descriptor/repformer_layer.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
-    List,
     Optional,
     Union,
 )
@@ -45,7 +44,7 @@ def get_residual(
     _mode: str = "norm",
     trainable: bool = True,
     precision: str = "float64",
-    seed: Optional[Union[int, List[int]]] = None,
+    seed: Optional[Union[int, list[int]]] = None,
 ) -> paddle.Tensor:
     r"""
     Get residual tensor for one update vector.
@@ -163,7 +162,7 @@ def __init__(
         smooth: bool = True,
         attnw_shift: float = 20.0,
         precision: str = "float64",
-        seed: Optional[Union[int, List[int]]] = None,
+        seed: Optional[Union[int, list[int]]] = None,
     ):
         """Return neighbor-wise multi-head self-attention maps, with gate mechanism."""
         super().__init__()
@@ -288,7 +287,7 @@ def __init__(
         input_dim: int,
         head_num: int,
         precision: str = "float64",
-        seed: Optional[Union[int, List[int]]] = None,
+        seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
         self.input_dim = input_dim
@@ -375,7 +374,7 @@ def __init__(
         input_dim: int,
         head_num: int,
         precision: str = "float64",
-        seed: Optional[Union[int, List[int]]] = None,
+        seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
         self.input_dim = input_dim
@@ -448,7 +447,7 @@ def __init__(
         smooth: bool = True,
         attnw_shift: float = 20.0,
         precision: str = "float64",
-        seed: Optional[Union[int, List[int]]] = None,
+        seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
         self.input_dim = input_dim
@@ -609,7 +608,7 @@ def __init__(
         precision: str = "float64",
         trainable_ln: bool = True,
         ln_eps: Optional[float] = 1e-5,
-        seed: Optional[Union[int, List[int]]] = None,
+        seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
         self.epsilon = 1e-4  # protection of 1./nnei
@@ -1070,10 +1069,10 @@ def forward(
         assert [nb, nloc] == g1.shape[:2]
         assert [nb, nloc, nnei] == h2.shape[:3]
 
-        g2_update: List[paddle.Tensor] = [g2]
-        h2_update: List[paddle.Tensor] = [h2]
-        g1_update: List[paddle.Tensor] = [g1]
-        g1_mlp: List[paddle.Tensor] = [g1]
+        g2_update: list[paddle.Tensor] = [g2]
+        h2_update: list[paddle.Tensor] = [h2]
+        g1_update: list[paddle.Tensor] = [g1]
+        g1_mlp: list[paddle.Tensor] = [g1]
 
         if cal_gg1:
             gg1 = _make_nei_g1(g1_ext, nlist)
@@ -1163,29 +1162,26 @@ def forward(
         g1_new = self.list_update(g1_update, "g1")
         return g1_new, g2_new, h2_new
 
-    # @paddle.jit.export
     def list_update_res_avg(
         self,
-        update_list: List[paddle.Tensor],
+        update_list: list[paddle.Tensor],
     ) -> paddle.Tensor:
         nitem = len(update_list)
         uu = update_list[0]
         for ii in range(1, nitem):
             uu = uu + update_list[ii]
         return uu / (float(nitem) ** 0.5)
 
-    # @paddle.jit.export
-    def list_update_res_incr(self, update_list: List[paddle.Tensor]) -> paddle.Tensor:
+    def list_update_res_incr(self, update_list: list[paddle.Tensor]) -> paddle.Tensor:
         nitem = len(update_list)
         uu = update_list[0]
         scale = 1.0 / (float(nitem - 1) ** 0.5) if nitem > 1 else 0.0
         for ii in range(1, nitem):
             uu = uu + scale * update_list[ii]
         return uu
 
-    # @paddle.jit.export
     def list_update_res_residual(
-        self, update_list: List[paddle.Tensor], update_name: str = "g1"
+        self, update_list: list[paddle.Tensor], update_name: str = "g1"
     ) -> paddle.Tensor:
         nitem = len(update_list)
         uu = update_list[0]
@@ -1203,9 +1199,8 @@ def list_update_res_residual(
             raise NotImplementedError
         return uu
 
-    # @paddle.jit.export
     def list_update(
-        self, update_list: List[paddle.Tensor], update_name: str = "g1"
+        self, update_list: list[paddle.Tensor], update_name: str = "g1"
     ) -> paddle.Tensor:
         if self.update_style == "res_avg":
             return self.list_update_res_avg(update_list)