test(dpa4): enable array-api-strict consistency tests

Author: njzjz

deepmd/dpmodel/descriptor/dpa4_nn/norm.py

@@ -546,7 +546,7 @@ def call(self, x: Any) -> Any:
         if x.ndim == 2:
             inv_rms = 1.0 / xp.sqrt(xp.mean(x * x, axis=-1, keepdims=True) + self.eps)
             x = x * inv_rms
-            x = x * xp_asarray_nodetach(xp, self.adam_scale[...], device=device)[0]
+            x = x * xp_asarray_nodetach(xp, self.adam_scale[...], device=device)[0, :]
             return xp.astype(x, in_dtype)
 
         inv_rms = 1.0 / xp.sqrt(xp.mean(x * x, axis=-1, keepdims=True) + self.eps)

source/tests/array_api_strict/common.py

@@ -59,8 +59,10 @@ def to_array_api_strict_array(array: np.ndarray | None) -> Any:
     f"{_PACKAGE_ROOT}.descriptor.dpa2",
     f"{_PACKAGE_ROOT}.descriptor.repflows",
     f"{_PACKAGE_ROOT}.descriptor.dpa3",
+    f"{_PACKAGE_ROOT}.descriptor.dpa4",
     f"{_PACKAGE_ROOT}.descriptor.hybrid",
     f"{_PACKAGE_ROOT}.fitting",
+    f"{_PACKAGE_ROOT}.fitting.dpa4_ener",
 )
 
 

source/tests/array_api_strict/descriptor/__init__.py

@@ -8,6 +8,9 @@
 from .dpa3 import (
     DescrptDPA3,
 )
+from .dpa4 import (
+    DescrptDPA4,
+)
 from .hybrid import (
     DescrptHybrid,
 )
@@ -31,6 +34,7 @@
     "DescrptDPA1",
     "DescrptDPA2",
     "DescrptDPA3",
+    "DescrptDPA4",
     "DescrptHybrid",
     "DescrptSeA",
     "DescrptSeAttenV2",

source/tests/array_api_strict/descriptor/dpa4.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP
+from deepmd.dpmodel.descriptor.dpa4_nn.activation import (
+    SwiGLU,
+)
+from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import (
+    GridProduct,
+)
+from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
+    BridgingSwitch,
+    C3CutoffEnvelope,
+    InnerClamp,
+)
+from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import (
+    WignerDCalculator,
+)
+
+from ..common import (
+    array_api_strict_module,
+    register_dpmodel_mapping,
+)
+from ..utils import exclude_mask as _strict_exclude_mask  # noqa: F401
+from ..utils import network as _strict_network  # noqa: F401
+from .base_descriptor import (
+    BaseDescriptor,
+)
+
+
+@BaseDescriptor.register("SeZM")
+@BaseDescriptor.register("sezm")
+@BaseDescriptor.register("DPA4")
+@BaseDescriptor.register("dpa4")
+@array_api_strict_module
+class DescrptDPA4(DescrptDPA4DP):
+    pass
+
+
+for _stateless_cls in (
+    BridgingSwitch,
+    C3CutoffEnvelope,
+    GridProduct,
+    InnerClamp,
+    SwiGLU,
+    WignerDCalculator,
+):
+    register_dpmodel_mapping(_stateless_cls, lambda v: v)

source/tests/array_api_strict/fitting/__init__.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from .dpa4_ener import (
+    SeZMEnergyFittingNet,
+)
 from .fitting import (
     DipoleFittingNet,
     DOSFittingNet,
@@ -13,4 +16,5 @@
     "EnergyFittingNet",
     "PolarFittingNet",
     "PropertyFittingNet",
+    "SeZMEnergyFittingNet",
 ]

source/tests/array_api_strict/fitting/dpa4_ener.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    ClassVar,
+)
+
+from deepmd.dpmodel.fitting.dpa4_ener import GLUFittingNet as GLUFittingNetDP
+from deepmd.dpmodel.fitting.dpa4_ener import (
+    SeZMEnergyFittingNet as SeZMEnergyFittingNetDP,
+)
+from deepmd.dpmodel.fitting.dpa4_ener import (
+    SeZMNetworkCollection as SeZMNetworkCollectionDP,
+)
+
+from ..common import (
+    array_api_strict_module,
+    register_dpmodel_mapping,
+)
+from ..utils import network as _strict_network  # noqa: F401
+
+
+@array_api_strict_module
+class GLUFittingNet(GLUFittingNetDP):
+    pass
+
+
+@array_api_strict_module
+class SeZMNetworkCollection(SeZMNetworkCollectionDP):
+    NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = {
+        "sezm_fitting_network": GLUFittingNet,
+    }
+
+
+@array_api_strict_module
+class SeZMEnergyFittingNet(SeZMEnergyFittingNetDP):
+    pass
+
+
+register_dpmodel_mapping(
+    GLUFittingNetDP,
+    lambda v: GLUFittingNet.deserialize(v.serialize()),
+)
+
+register_dpmodel_mapping(
+    SeZMNetworkCollectionDP,
+    lambda v: SeZMNetworkCollection.deserialize(v.serialize()),
+)

source/tests/consistent/descriptor/test_dpa4.py

@@ -19,6 +19,7 @@
 )
 
 from ..common import (
+    INSTALLED_ARRAY_API_STRICT,
     INSTALLED_PT,
     INSTALLED_PT_EXPT,
     CommonTest,
@@ -28,14 +29,40 @@
     DescriptorTest,
 )
 
-if INSTALLED_PT:
-    from deepmd.pt.model.descriptor.sezm import DescrptSeZM as DescrptDPA4PT
-else:
-    DescrptDPA4PT = None
-if INSTALLED_PT_EXPT:
-    from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt
+DescrptDPA4PT = None
+DescrptDPA4PTExpt = None
+
+
+def _get_descrpt_dpa4_pt() -> type | None:
+    global DescrptDPA4PT
+    if not INSTALLED_PT:
+        return None
+    if DescrptDPA4PT is None:
+        from deepmd.pt.model.descriptor.sezm import (
+            DescrptSeZM,
+        )
+
+        DescrptDPA4PT = DescrptSeZM
+    return DescrptDPA4PT
+
+
+def _get_descrpt_dpa4_pt_expt() -> type | None:
+    global DescrptDPA4PTExpt
+    if not INSTALLED_PT_EXPT:
+        return None
+    if DescrptDPA4PTExpt is None:
+        from deepmd.pt_expt.descriptor.dpa4 import (
+            DescrptDPA4,
+        )
+
+        DescrptDPA4PTExpt = DescrptDPA4
+    return DescrptDPA4PTExpt
+
+
+if INSTALLED_ARRAY_API_STRICT:
+    from ...array_api_strict.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4Strict
 else:
-    DescrptDPA4PTExpt = None
+    DescrptDPA4Strict = None
 
 # not implemented
 DescrptDPA4TF = None
@@ -146,22 +173,31 @@ def data(self) -> dict:
 
     @property
     def skip_pt(self) -> bool:
-        return CommonTest.skip_pt
+        return CommonTest.skip_pt or _get_descrpt_dpa4_pt() is None
+
+    @property
+    def skip_pt_expt(self) -> bool:
+        return not INSTALLED_PT_EXPT or _get_descrpt_dpa4_pt_expt() is None
+
+    @property
+    def pt_class(self) -> type | None:
+        return _get_descrpt_dpa4_pt()
+
+    @property
+    def pt_expt_class(self) -> type | None:
+        return _get_descrpt_dpa4_pt_expt()
 
     skip_dp = False
     skip_tf = True
     skip_jax = True
     skip_pd = True
-    skip_pt_expt = not INSTALLED_PT_EXPT
-    skip_array_api_strict = True
+    skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT
 
     tf_class = DescrptDPA4TF
     dp_class = DescrptDPA4DP
-    pt_class = DescrptDPA4PT
-    pt_expt_class = DescrptDPA4PTExpt
     jax_class = None
     pd_class = None
-    array_api_strict_class = None
+    array_api_strict_class = DescrptDPA4Strict
     args: ClassVar[list] = [
         *descrpt_se_zm_args(),
         Argument("ntypes", int, optional=False),
@@ -234,6 +270,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any:
             mixed_types=True,
         )
 
+    def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
+        return self.eval_array_api_strict_descriptor(
+            array_api_strict_obj,
+            self.natoms,
+            self.coords,
+            self.atype,
+            self.box,
+            mixed_types=True,
+        )
+
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
         return (ret[0],)