fix ut

Author: iProzd

deepmd/jax/descriptor/dpa3.py

@@ -23,6 +23,9 @@
     flax_version,
     nnx,
 )
+from deepmd.jax.utils.network import (
+    NativeLayer,
+)
 from deepmd.jax.utils.type_embed import (
     TypeEmbedNet,
 )
@@ -40,8 +43,12 @@ def __setattr__(self, name: str, value: Any) -> None:
                 value = nnx.data(value)
         elif name in {"repflows"}:
             value = DescrptBlockRepflows.deserialize(value.serialize())
-        elif name in {"type_embedding"}:
-            value = TypeEmbedNet.deserialize(value.serialize())
+        elif name in {"type_embedding", "chg_embedding", "spin_embedding"}:
+            if value is not None:
+                value = TypeEmbedNet.deserialize(value.serialize())
+        elif name in {"mix_cs_mlp"}:
+            if value is not None:
+                value = NativeLayer.deserialize(value.serialize())
         else:
             pass
         return super().__setattr__(name, value)

source/tests/array_api_strict/descriptor/dpa3.py

@@ -8,6 +8,9 @@
 from ..common import (
     to_array_api_strict_array,
 )
+from ..utils.network import (
+    NativeLayer,
+)
 from ..utils.type_embed import (
     TypeEmbedNet,
 )
@@ -26,8 +29,12 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_array_api_strict_array(value)
         elif name in {"repflows"}:
             value = DescrptBlockRepflows.deserialize(value.serialize())
-        elif name in {"type_embedding"}:
-            value = TypeEmbedNet.deserialize(value.serialize())
+        elif name in {"type_embedding", "chg_embedding", "spin_embedding"}:
+            if value is not None:
+                value = TypeEmbedNet.deserialize(value.serialize())
+        elif name in {"mix_cs_mlp"}:
+            if value is not None:
+                value = NativeLayer.deserialize(value.serialize())
         else:
             pass
         return super().__setattr__(name, value)