style(tf): add type annotations (#4945)

Addressed review feedback to replace remaining `Any` type annotations
with specific types across training, fitting, and model modules.

## Changes

### Learning Rate Types
- `EnerFitting.get_loss()` and `StandardModel.get_loss()`: `lr: Any` →
`lr: LearningRateExp`

### Trainer Method Signatures
- `build()`: `data: dict` → `data: DeepmdDataSystem | None`
- `_build_loss()`: `-> Any` → `-> tuple[None, None] | tuple[tf.Tensor,
dict[str, tf.Tensor]]`
- `_build_network()`: `data: dict` → `data: DeepmdDataSystem`
- `get_global_step()`: `-> Any` → `-> int`

### Example

```python
# Before
def get_loss(self, loss: dict, lr: Any) -> Loss:
    ...

# After
def get_loss(self, loss: dict, lr: LearningRateExp) -> Loss:
    ...
```

The `_build_loss()` return type now accurately reflects that it returns
either `(None, None)` when `stop_batch == 0`, or `(tf.Tensor, dict[str,
tf.Tensor])` from the loss build.

<!-- START COPILOT CODING AGENT TIPS -->
---

✨ Let Copilot coding agent [set things up for
you](https://github.com/deepmodeling/deepmd-kit/issues/new?title=✨+Set+up+Copilot+instructions&body=Configure%20instructions%20for%20this%20repository%20as%20documented%20in%20%5BBest%20practices%20for%20Copilot%20coding%20agent%20in%20your%20repository%5D%28https://gh.io/copilot-coding-agent-tips%29%2E%0A%0A%3COnboard%20this%20repo%3E&assignees=copilot)
— coding agent works faster and does higher quality work when set up for
your repo.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

deepmd/tf/cluster/local.py

@@ -14,7 +14,7 @@
 __all__ = ["get_gpus", "get_resource"]
 
 
-def get_gpus():
+def get_gpus() -> list[int] | None:
     """Get available IDs of GPU cards at local.
     These IDs are valid when used as the TensorFlow device ID.
 

deepmd/tf/common.py

@@ -14,6 +14,7 @@
     Union,
 )
 
+import numpy as np
 import tensorflow
 from packaging.version import (
     Version,
@@ -118,7 +119,7 @@ def gelu_tf(x: tf.Tensor) -> tf.Tensor:
     https://arxiv.org/abs/1606.08415
     """
 
-    def gelu_wrapper(x):
+    def gelu_wrapper(x: tf.Tensor) -> tf.Tensor:
         try:
             return tensorflow.nn.gelu(x, approximate=True)
         except AttributeError:
@@ -146,16 +147,14 @@ def silu(x: tf.Tensor) -> tf.Tensor:
     return x * tf.sigmoid(x)
 
 
-def get_silut(activation_function: str = "silut"):
-    import numpy as np
-
-    def sigmoid(x):
+def get_silut(activation_function: str = "silut") -> Callable[[tf.Tensor], tf.Tensor]:
+    def sigmoid(x: float | np.ndarray) -> float | np.ndarray:
         return 1 / (1 + np.exp(-x))
 
-    def silu(x):
+    def silu(x: float | np.ndarray) -> float | np.ndarray:
         return x * sigmoid(x)
 
-    def silu_grad(x):
+    def silu_grad(x: float | np.ndarray) -> float | np.ndarray:
         sig = sigmoid(x)
         return sig + x * sig * (1 - sig)
 
@@ -233,7 +232,7 @@ def get_activation_func(
     return ACTIVATION_FN_DICT[activation_fn.lower()]
 
 
-def get_precision(precision: "_PRECISION") -> Any:
+def get_precision(precision: "_PRECISION") -> tf.DType:
     """Convert str to TF DType constant.
 
     Parameters
@@ -319,7 +318,7 @@ def cast_precision(func: Callable) -> Callable:
     """
 
     @wraps(func)
-    def wrapper(self, *args, **kwargs):
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
         # only convert tensors
         returned_tensor = func(
             self,

deepmd/tf/descriptor/descriptor.py

@@ -3,11 +3,17 @@
     abstractmethod,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from typing_extensions import (
+        Self,
+    )
+
 from deepmd.common import (
     j_get_type,
 )
@@ -48,7 +54,7 @@ class Descriptor(PluginVariant, make_plugin_registry("descriptor")):
     that can be called by other classes.
     """
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Self":
         if cls is Descriptor:
             cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
         return super().__new__(cls)
@@ -132,7 +138,7 @@ def compute_input_stats(
         natoms_vec: list[np.ndarray],
         mesh: list[np.ndarray],
         input_dict: dict[str, list[np.ndarray]],
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Compute the statisitcs (avg and std) of the training data. The input will be
         normalized by the statistics.

deepmd/tf/descriptor/hybrid.py

@@ -46,7 +46,7 @@ def __init__(
         list: list[Descriptor | dict[str, Any]],
         ntypes: int | None = None,
         spin: Spin | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Constructor."""
         # warning: list is conflict with built-in list
@@ -142,7 +142,7 @@ def compute_input_stats(
         input_dict: dict,
         mixed_type: bool = False,
         real_natoms_vec: list | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
 
@@ -182,7 +182,7 @@ def compute_input_stats(
                 **kwargs,
             )
 
-    def merge_input_stats(self, stat_dict) -> None:
+    def merge_input_stats(self, stat_dict: dict[str, float]) -> None:
         """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
 
         Parameters

deepmd/tf/descriptor/loc_frame.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
 
 import numpy as np
 
@@ -58,7 +61,7 @@ def __init__(
         sel_a: list[int],
         sel_r: list[int],
         axis_rule: list[int],
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Constructor."""
         self.sel_a = sel_a
@@ -159,7 +162,7 @@ def compute_input_stats(
         natoms_vec: list,
         mesh: list,
         input_dict: dict,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
 
@@ -372,8 +375,13 @@ def prod_force_virial(
         return force, virial, atom_virial
 
     def _compute_dstats_sys_nonsmth(
-        self, data_coord, data_box, data_atype, natoms_vec, mesh
-    ):
+        self,
+        data_coord: np.ndarray,
+        data_box: np.ndarray,
+        data_atype: np.ndarray,
+        natoms_vec: np.ndarray,
+        mesh: np.ndarray,
+    ) -> tuple[Any, Any]:
         dd_all = run_sess(
             self.sub_sess,
             self.stat_descrpt,
@@ -405,7 +413,7 @@ def _compute_dstats_sys_nonsmth(
             sysv2.append(sumv2)
         return sysv, sysv2, sysn
 
-    def _compute_std(self, sumv2, sumv, sumn):
+    def _compute_std(self, sumv2: float, sumv: float, sumn: float) -> float:
         return np.sqrt(sumv2 / sumn - np.multiply(sumv / sumn, sumv / sumn))
 
     def init_variables(

deepmd/tf/descriptor/se_a.py

@@ -1,5 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from collections.abc import (
+    Callable,
+)
+from typing import (
+    Any,
+)
+
 import numpy as np
 
 from deepmd.dpmodel.utils.env_mat import (
@@ -180,7 +187,7 @@ def __init__(
         tebd_input_mode: str = "concat",
         type_map: list[str] | None = None,  # to be compat with input
         env_protection: float = 0.0,  # not implement!!
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Constructor."""
         if rcut < rcut_smth:
@@ -346,7 +353,7 @@ def compute_input_stats(
         natoms_vec: list,
         mesh: list,
         input_dict: dict,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
 
@@ -393,7 +400,7 @@ def compute_input_stats(
             }
             self.merge_input_stats(stat_dict)
 
-    def merge_input_stats(self, stat_dict) -> None:
+    def merge_input_stats(self, stat_dict: dict[str, Any]) -> None:
         """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
 
         Parameters
@@ -747,8 +754,15 @@ def prod_force_virial(
         return force, virial, atom_virial
 
     def _pass_filter(
-        self, inputs, atype, natoms, input_dict, reuse=None, suffix="", trainable=True
-    ):
+        self,
+        inputs: tf.Tensor,
+        atype: tf.Tensor,
+        natoms: tf.Tensor,
+        input_dict: dict,
+        reuse: bool | None = None,
+        suffix: str = "",
+        trainable: bool = True,
+    ) -> tuple[tf.Tensor, tf.Tensor]:
         if input_dict is not None:
             type_embedding = input_dict.get("type_embedding", None)
             if type_embedding is not None:
@@ -834,8 +848,13 @@ def _pass_filter(
         return output, output_qmat
 
     def _compute_dstats_sys_smth(
-        self, data_coord, data_box, data_atype, natoms_vec, mesh
-    ):
+        self,
+        data_coord: np.ndarray,
+        data_box: np.ndarray,
+        data_atype: np.ndarray,
+        natoms_vec: np.ndarray,
+        mesh: np.ndarray,
+    ) -> tuple[list[float], list[float], list[float], list[float], list[int]]:
         dd_all = run_sess(
             self.sub_sess,
             self.stat_descrpt,
@@ -876,7 +895,7 @@ def _compute_dstats_sys_smth(
             sysa2.append(suma2)
         return sysr, sysr2, sysa, sysa2, sysn
 
-    def _compute_std(self, sumv2, sumv, sumn):
+    def _compute_std(self, sumv2: float, sumv: float, sumn: int) -> float:
         if sumn == 0:
             return 1.0 / self.rcut_r
         val = np.sqrt(sumv2 / sumn - np.multiply(sumv / sumn, sumv / sumn))
@@ -886,11 +905,11 @@ def _compute_std(self, sumv2, sumv, sumn):
 
     def _concat_type_embedding(
         self,
-        xyz_scatter,
-        nframes,
-        natoms,
-        type_embedding,
-    ):
+        xyz_scatter: tf.Tensor,
+        nframes: tf.Tensor,
+        natoms: tf.Tensor,
+        type_embedding: tf.Tensor,
+    ) -> tf.Tensor:
         """Concatenate `type_embedding` of neighbors and `xyz_scatter`.
         If not self.type_one_side, concatenate `type_embedding` of center atoms as well.
 
@@ -939,21 +958,21 @@ def _concat_type_embedding(
 
     def _filter_lower(
         self,
-        type_i,
-        type_input,
-        start_index,
-        incrs_index,
-        inputs,
-        nframes,
-        natoms,
-        type_embedding=None,
-        is_exclude=False,
-        activation_fn=None,
-        bavg=0.0,
-        stddev=1.0,
-        trainable=True,
-        suffix="",
-    ):
+        type_i: int,
+        type_input: int,
+        start_index: int,
+        incrs_index: int,
+        inputs: tf.Tensor,
+        nframes: tf.Tensor,
+        natoms: tf.Tensor,
+        type_embedding: tf.Tensor | None = None,
+        is_exclude: bool = False,
+        activation_fn: Callable[[tf.Tensor], tf.Tensor] | None = None,
+        bavg: float = 0.0,
+        stddev: float = 1.0,
+        trainable: bool = True,
+        suffix: str = "",
+    ) -> tf.Tensor:
         """Input env matrix, returns R.G."""
         outputs_size = [1, *self.filter_neuron]
         # cut-out inputs
@@ -1159,17 +1178,17 @@ def _filter_lower(
     @cast_precision
     def _filter(
         self,
-        inputs,
-        type_input,
-        natoms,
-        type_embedding=None,
-        activation_fn=tf.nn.tanh,
-        stddev=1.0,
-        bavg=0.0,
-        name="linear",
-        reuse=None,
-        trainable=True,
-    ):
+        inputs: tf.Tensor,
+        type_input: int,
+        natoms: tf.Tensor,
+        type_embedding: tf.Tensor | None = None,
+        activation_fn: Callable[[tf.Tensor], tf.Tensor] | None = tf.nn.tanh,
+        stddev: float = 1.0,
+        bavg: float = 0.0,
+        name: str = "linear",
+        reuse: bool | None = None,
+        trainable: bool = True,
+    ) -> tf.Tensor:
         nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0]
         # natom x (nei x 4)
         shape = inputs.get_shape().as_list()
@@ -1363,7 +1382,7 @@ def explicit_ntypes(self) -> bool:
         return False
 
     @classmethod
-    def deserialize(cls, data: dict, suffix: str = ""):
+    def deserialize(cls, data: dict, suffix: str = "") -> "DescrptSeA":
         """Deserialize the model.
 
         Parameters