feat(jax): enable ANN rule and add comprehensive type hints to JAX backend

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/jax/utils/serialization.py

@@ -2,9 +2,6 @@
 from pathlib import (
     Path,
 )
-from typing import (
-    Any,
-)
 
 import numpy as np
 import orbax.checkpoint as ocp
@@ -58,7 +55,7 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
 
         def exported_whether_do_atomic_virial(
             do_atomic_virial: bool, has_ghost_atoms: bool
-        ) -> Any:
+        ) -> jax_export.Exported:
             def call_lower_with_fixed_do_atomic_virial(
                 coord: jnp.ndarray,
                 atype: jnp.ndarray,

pyproject.toml

@@ -425,7 +425,6 @@ runtime-evaluated-base-classes = ["torch.nn.Module"]
 "deepmd/tf/**" = ["TID253", "ANN"]
 "deepmd/pt/**" = ["TID253"]
 "deepmd/jax/**" = ["TID253"]
-"deepmd/jax/jax2tf/**" = ["TID253", "ANN"]
 "deepmd/pd/**" = ["TID253", "ANN"]
 
 "source/**" = ["ANN"]