diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/norm.py b/deepmd/dpmodel/descriptor/dpa4_nn/norm.py index f436f8c3dc..c6492bd561 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/norm.py +++ b/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) diff --git a/source/tests/array_api_strict/common.py b/source/tests/array_api_strict/common.py index 546e111dd2..b1032c4bd9 100644 --- a/source/tests/array_api_strict/common.py +++ b/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", ) diff --git a/source/tests/array_api_strict/descriptor/__init__.py b/source/tests/array_api_strict/descriptor/__init__.py index 1bbefbea6f..b9235aa8b5 100644 --- a/source/tests/array_api_strict/descriptor/__init__.py +++ b/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", diff --git a/source/tests/array_api_strict/descriptor/dpa4.py b/source/tests/array_api_strict/descriptor/dpa4.py new file mode 100644 index 0000000000..02630d6f6a --- /dev/null +++ b/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) diff --git a/source/tests/array_api_strict/fitting/__init__.py b/source/tests/array_api_strict/fitting/__init__.py index 2041f600ea..81aee58175 100644 --- a/source/tests/array_api_strict/fitting/__init__.py +++ b/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", ] diff --git a/source/tests/array_api_strict/fitting/dpa4_ener.py b/source/tests/array_api_strict/fitting/dpa4_ener.py new file mode 100644 index 0000000000..d7211fc4e5 --- /dev/null +++ b/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 diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index e6f3216bd4..1cfbf2f9a1 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/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,7 +158,7 @@ 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 @@ -161,7 +166,7 @@ def skip_pt(self) -> bool: 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],) diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index ee10b54343..3d007a9959 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/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,7 +89,7 @@ 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 @@ -85,7 +97,7 @@ def skip_pt(self) -> bool: 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,)