fix(jax): improve JAX modules' names (#5214)

Use `wraps` to keep the modules' names, so they won't be `FlaxModule`,
which cannot be regonized. I realized it when implementing #5213.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Improved type annotations and type-safety for module wrapping to aid
correctness and tooling.
* Preserved original module metadata when creating wrapped modules,
improving introspection and debugging.
* Updated docstrings and signatures to reflect the new typing and
metadata behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-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: njzjz

deepmd/jax/common.py

@@ -1,14 +1,15 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from functools import (
+    wraps,
+)
 from typing import (
     Any,
+    TypeVar,
     overload,
 )
 
 import numpy as np
 
-from deepmd.dpmodel.common import (
-    NativeOP,
-)
 from deepmd.jax.env import (
     jnp,
     nnx,
@@ -41,19 +42,22 @@ def to_jax_array(array: np.ndarray | None) -> jnp.ndarray | None:
     return jnp.array(array)
 
 
+T = TypeVar("T")
+
+
 def flax_module(
-    module: NativeOP,
-) -> nnx.Module:
+    module: type[T],
+) -> type[T]:  # runtime: actually returns type[T & nnx.Module]
     """Convert a NativeOP to a Flax module.
 
     Parameters
     ----------
-    module : NativeOP
+    module : type[NativeOP]
         The NativeOP to convert.
 
     Returns
     -------
-    flax.nnx.Module
+    type[flax.nnx.Module]
         The Flax module.
 
     Examples
@@ -72,6 +76,7 @@ class MixedMetaClass(*metas):
         def __call__(self, *args: Any, **kwargs: Any) -> Any:
             return type(nnx.Module).__call__(self, *args, **kwargs)
 
+    @wraps(module, updated=())
     class FlaxModule(module, nnx.Module, metaclass=MixedMetaClass):
         def __init_subclass__(cls, **kwargs: Any) -> None:
             return super().__init_subclass__(**kwargs)