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

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,51 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from importlib import (
+    import_module,
+)
+
+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 .base_descriptor import (
+    BaseDescriptor,
+)
+
+import_module("..utils.exclude_mask", __package__)
+import_module("..utils.network", __package__)
+
+
+@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,38 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from importlib import (
+    import_module,
+)
+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,
+)
+
+import_module("..utils.network", __package__)
+
+
+@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

source/tests/consistent/descriptor/test_dpa4.py

@@ -19,6 +19,7 @@
 )
 
 from ..common import (
+    INSTALLED_ARRAY_API_STRICT,
     INSTALLED_PT,
     INSTALLED_PT_EXPT,
     CommonTest,
@@ -36,6 +37,10 @@
     from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt
 else:
     DescrptDPA4PTExpt = None
+if INSTALLED_ARRAY_API_STRICT:
+    from ...array_api_strict.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4Strict
+else:
+    DescrptDPA4Strict = None
 
 # not implemented
 DescrptDPA4TF = None
@@ -153,15 +158,15 @@ def skip_pt(self) -> bool:
     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 +239,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],)
 

source/tests/consistent/fitting/test_dpa4_ener.py

@@ -6,6 +6,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
 from deepmd.dpmodel.fitting.dpa4_ener import SeZMEnergyFittingNet as SeZMEnerFittingDP
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
@@ -15,6 +18,7 @@
 )
 
 from ..common import (
+    INSTALLED_ARRAY_API_STRICT,
     INSTALLED_PT,
     INSTALLED_PT_EXPT,
     CommonTest,
@@ -40,6 +44,14 @@
     from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE
 else:
     SeZMEnerFittingPTExpt = None
+if INSTALLED_ARRAY_API_STRICT:
+    import array_api_strict
+
+    from ...array_api_strict.fitting.dpa4_ener import (
+        SeZMEnergyFittingNet as SeZMEnerFittingStrict,
+    )
+else:
+    SeZMEnerFittingStrict = None
 
 # not implemented
 SeZMEnerFittingTF = None
@@ -77,15 +89,15 @@ def skip_pt(self) -> bool:
     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 = SeZMEnerFittingTF
     dp_class = SeZMEnerFittingDP
     pt_class = SeZMEnerFittingPT
     pt_expt_class = SeZMEnerFittingPTExpt
     jax_class = None
     pd_class = None
-    array_api_strict_class = None
+    array_api_strict_class = SeZMEnerFittingStrict
     args = fitting_sezm_ener()
 
     def setUp(self) -> None:
@@ -138,6 +150,14 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any:
             .numpy()
         )
 
+    def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
+        return to_numpy_array(
+            array_api_strict_obj(
+                array_api_strict.asarray(self.inputs),
+                array_api_strict.asarray(self.atype.reshape(1, -1)),
+            )["energy"]
+        )
+
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
         return (ret,)