test(dpa4): simplify backend imports

Restore the established eager conditional-import pattern for the PyTorch backends while retaining the array-api-strict coverage.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning-Effort: xhigh

Author: njzjz-bot

source/tests/consistent/descriptor/test_dpa4.py

@@ -29,36 +29,14 @@
     DescriptorTest,
 )
 
-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_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
+else:
+    DescrptDPA4PTExpt = None
 if INSTALLED_ARRAY_API_STRICT:
     from ...array_api_strict.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4Strict
 else:
@@ -173,28 +151,19 @@ def data(self) -> dict:
 
     @property
     def skip_pt(self) -> bool:
-        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()
+        return CommonTest.skip_pt
 
     skip_dp = False
     skip_tf = True
     skip_jax = True
     skip_pd = True
+    skip_pt_expt = not INSTALLED_PT_EXPT
     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 = DescrptDPA4Strict

source/tests/consistent/fitting/test_dpa4_ener.py

@@ -28,53 +28,22 @@
     FittingTest,
 )
 
-torch = None
-PT_DEVICE = None
-PT_EXPT_DEVICE = None
-SeZMEnerFittingPT = None
-SeZMEnerFittingPTExpt = None
-
-
-def _get_sezm_ener_fitting_pt() -> type | None:
-    global PT_DEVICE, SeZMEnerFittingPT, torch
-    if not INSTALLED_PT:
-        return None
-    if SeZMEnerFittingPT is None:
-        import torch as torch_module
-
-        from deepmd.pt.model.task.sezm_ener import (
-            SeZMEnergyFittingNet,
-        )
-        from deepmd.pt.utils.env import (
-            DEVICE,
-        )
-
-        torch = torch_module
-        PT_DEVICE = DEVICE
-        SeZMEnerFittingPT = SeZMEnergyFittingNet
-    return SeZMEnerFittingPT
-
-
-def _get_sezm_ener_fitting_pt_expt() -> type | None:
-    global PT_EXPT_DEVICE, SeZMEnerFittingPTExpt, torch
-    if not INSTALLED_PT_EXPT:
-        return None
-    if SeZMEnerFittingPTExpt is None:
-        import torch as torch_module
-
-        from deepmd.pt_expt.fitting.dpa4_ener import (
-            SeZMEnergyFittingNet,
-        )
-        from deepmd.pt_expt.utils.env import (
-            DEVICE,
-        )
-
-        torch = torch_module
-        PT_EXPT_DEVICE = DEVICE
-        SeZMEnerFittingPTExpt = SeZMEnergyFittingNet
-    return SeZMEnerFittingPTExpt
+if INSTALLED_PT:
+    import torch
 
+    from deepmd.pt.model.task.sezm_ener import SeZMEnergyFittingNet as SeZMEnerFittingPT
+    from deepmd.pt.utils.env import DEVICE as PT_DEVICE
+else:
+    SeZMEnerFittingPT = None
+if INSTALLED_PT_EXPT:
+    import torch
 
+    from deepmd.pt_expt.fitting.dpa4_ener import (
+        SeZMEnergyFittingNet as SeZMEnerFittingPTExpt,
+    )
+    from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE
+else:
+    SeZMEnerFittingPTExpt = None
 if INSTALLED_ARRAY_API_STRICT:
     import array_api_strict
 
@@ -113,28 +82,19 @@ def data(self) -> dict:
 
     @property
     def skip_pt(self) -> bool:
-        return CommonTest.skip_pt or _get_sezm_ener_fitting_pt() is None
-
-    @property
-    def skip_pt_expt(self) -> bool:
-        return not INSTALLED_PT_EXPT or _get_sezm_ener_fitting_pt_expt() is None
-
-    @property
-    def pt_class(self) -> type | None:
-        return _get_sezm_ener_fitting_pt()
-
-    @property
-    def pt_expt_class(self) -> type | None:
-        return _get_sezm_ener_fitting_pt_expt()
+        return CommonTest.skip_pt
 
     skip_dp = False
     skip_tf = True
     skip_jax = True
     skip_pd = True
+    skip_pt_expt = not INSTALLED_PT_EXPT
     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 = SeZMEnerFittingStrict