fix(jax): improve JAX modules' names

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

Author: njzjz

deepmd/jax/common.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from functools import (
+    wraps,
+)
 from typing import (
     Any,
     overload,
@@ -42,18 +45,18 @@ def to_jax_array(array: np.ndarray | None) -> jnp.ndarray | None:
 
 
 def flax_module(
-    module: NativeOP,
-) -> nnx.Module:
+    module: type[NativeOP],
+) -> type[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 +75,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)