test(pt_expt): port silut activation + repformers accessors from deepmodeling#5393

Add silut/custom_silu support to _torch_activation using native torch
ops (sigmoid, tanh, where) so the custom silu stays traceable by
make_fx / torch.export. Cross-backend consistency tests cover multiple
thresholds across the silu/tanh branches, and a pt_expt unit file
exercises default/custom threshold, gradient flow, make_fx, and
torch.export.

Also port DescrptBlockRepformers accessor tests (get_rcut_smth,
get_env_protection). The underlying accessor methods already exist on
this branch; these tests guard against regressions.

Author: Han Wang

deepmd/pt_expt/utils/network.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import math
 from typing import (
     Any,
     ClassVar,
@@ -182,6 +183,14 @@ def _torch_activation(x: torch.Tensor, name: str) -> torch.Tensor:
         return torch.sigmoid(x)
     elif name == "silu":
         return torch.nn.functional.silu(x)
+    elif name.startswith("silut") or name.startswith("custom_silu"):
+        threshold = float(name.split(":")[-1]) if ":" in name else 3.0
+        sig_t = 1.0 / (1.0 + math.exp(-threshold))
+        slope = sig_t + threshold * sig_t * (1.0 - sig_t)
+        const = threshold * sig_t
+        silu = x * torch.sigmoid(x)
+        tanh_branch = torch.tanh(slope * (x - threshold)) + const
+        return torch.where(x < threshold, silu, tanh_branch)
     elif name in ("none", "linear"):
         return x
     else:

source/tests/common/dpmodel/test_descriptor_dpa2.py

@@ -10,6 +10,9 @@
     RepformerArgs,
     RepinitArgs,
 )
+from deepmd.dpmodel.descriptor.repformers import (
+    DescrptBlockRepformers,
+)
 
 from ...seed import (
     GLOBAL_SEED,
@@ -69,3 +72,36 @@ def test_self_consistency(
         for ii in [0, 1, 2, 3, 4]:
             np.testing.assert_equal(mm0[ii].shape, desired_shape[ii])
             np.testing.assert_allclose(mm0[ii], mm1[ii])
+
+
+class TestDescrptBlockRepformersAccessors(unittest.TestCase):
+    def test_get_rcut_smth(self) -> None:
+        block = DescrptBlockRepformers(
+            rcut=6.0,
+            rcut_smth=5.0,
+            sel=40,
+            ntypes=2,
+            nlayers=3,
+        )
+        self.assertEqual(block.get_rcut_smth(), 5.0)
+
+    def test_get_env_protection(self) -> None:
+        block = DescrptBlockRepformers(
+            rcut=6.0,
+            rcut_smth=5.0,
+            sel=40,
+            ntypes=2,
+            nlayers=3,
+            env_protection=1.0,
+        )
+        self.assertEqual(block.get_env_protection(), 1.0)
+
+    def test_get_env_protection_default(self) -> None:
+        block = DescrptBlockRepformers(
+            rcut=6.0,
+            rcut_smth=5.0,
+            sel=40,
+            ntypes=2,
+            nlayers=3,
+        )
+        self.assertEqual(block.get_env_protection(), 0.0)

source/tests/consistent/test_activation.py

@@ -19,6 +19,7 @@
     INSTALLED_JAX,
     INSTALLED_PD,
     INSTALLED_PT,
+    INSTALLED_PT_EXPT,
     INSTALLED_TF,
     parameterized,
 )
@@ -29,6 +30,13 @@
     from deepmd.pt.utils.utils import (
         to_torch_tensor,
     )
+if INSTALLED_PT_EXPT:
+    import torch
+
+    from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE
+    from deepmd.pt_expt.utils.network import (
+        _torch_activation,
+    )
 if INSTALLED_TF:
     from deepmd.tf.common import get_activation_func as get_activation_fn_tf
     from deepmd.tf.env import (
@@ -98,3 +106,54 @@ def test_pd_consistent_with_ref(self):
                 ActivationFn_pd(self.activation)(to_paddle_tensor(self.random_input))
             )
             np.testing.assert_allclose(self.ref, test, atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_PT_EXPT, "PyTorch Exportable is not installed")
+    def test_pt_expt_consistent_with_ref(self) -> None:
+        if INSTALLED_PT_EXPT:
+            x = torch.tensor(
+                self.random_input, dtype=torch.float64, device=PT_EXPT_DEVICE
+            )
+            test = _torch_activation(x, self.activation).detach().numpy()
+            np.testing.assert_allclose(self.ref, test, atol=1e-10)
+
+
+@parameterized(
+    (
+        "silut",  # default threshold 3.0
+        "silut:3.0",  # explicit threshold 3.0
+        "silut:10.0",  # large threshold
+        "custom_silu:5.0",  # alias
+    ),
+)
+class TestSilutVariantsConsistent(unittest.TestCase):
+    """Cross-backend consistency for silut with different thresholds."""
+
+    def setUp(self) -> None:
+        (self.activation,) = self.param
+        # Parse threshold to build input that covers both branches
+        threshold = (
+            float(self.activation.split(":")[-1]) if ":" in self.activation else 3.0
+        )
+        rng = np.random.default_rng(GLOBAL_SEED)
+        # Values below threshold (silu branch) and above threshold (tanh branch)
+        below = rng.uniform(-threshold - 5, threshold - 0.1, size=(5, 10))
+        above = rng.uniform(threshold + 0.1, threshold + 20, size=(5, 10))
+        self.random_input = np.concatenate([below, above], axis=0)
+        self.ref = get_activation_fn_dp(self.activation)(self.random_input)
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_pt_consistent_with_ref(self) -> None:
+        if INSTALLED_PT:
+            test = torch_to_numpy(
+                ActivationFn_pt(self.activation)(to_torch_tensor(self.random_input))
+            )
+            np.testing.assert_allclose(self.ref, test, atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_PT_EXPT, "PyTorch Exportable is not installed")
+    def test_pt_expt_consistent_with_ref(self) -> None:
+        if INSTALLED_PT_EXPT:
+            x = torch.tensor(
+                self.random_input, dtype=torch.float64, device=PT_EXPT_DEVICE
+            )
+            test = _torch_activation(x, self.activation).detach().numpy()
+            np.testing.assert_allclose(self.ref, test, atol=1e-10)

source/tests/pt_expt/utils/test_activation.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+import torch
+from torch.fx.experimental.proxy_tensor import (
+    make_fx,
+)
+
+from deepmd.dpmodel.utils.network import (
+    get_activation_fn,
+)
+from deepmd.pt_expt.utils.network import (
+    _torch_activation,
+)
+
+
+class TestSilutActivation:
+    """Tests for silut activation in _torch_activation."""
+
+    def setup_method(self) -> None:
+        # x values spanning both branches: below threshold and above
+        self.x_np = np.array(
+            [-5.0, -1.0, 0.0, 1.0, 2.5, 3.0, 5.0, 10.0, 15.0, 20.0],
+            dtype=np.float64,
+        )
+        self.x_torch = torch.tensor(self.x_np, dtype=torch.float64)
+
+    def test_silut_with_threshold(self) -> None:
+        """silut:10.0 matches dpmodel numerically."""
+        result = _torch_activation(self.x_torch, "silut:10.0")
+        dp_fn = get_activation_fn("silut:10.0")
+        expected = dp_fn(self.x_np)
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected, rtol=1e-12, atol=1e-12
+        )
+
+    def test_silut_default_threshold(self) -> None:
+        """Silut without parameter uses default threshold 3.0."""
+        result = _torch_activation(self.x_torch, "silut")
+        dp_fn = get_activation_fn("silut")
+        expected = dp_fn(self.x_np)
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected, rtol=1e-12, atol=1e-12
+        )
+
+    def test_silut_custom_silu_alias(self) -> None:
+        """custom_silu:5.0 is an alias for silut:5.0."""
+        result = _torch_activation(self.x_torch, "custom_silu:5.0")
+        dp_fn = get_activation_fn("custom_silu:5.0")
+        expected = dp_fn(self.x_np)
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected, rtol=1e-12, atol=1e-12
+        )
+
+    def test_silut_gradient(self) -> None:
+        """Gradient flows through both branches of silut."""
+        x = self.x_torch.clone().requires_grad_(True)
+        y = _torch_activation(x, "silut:3.0")
+        loss = y.sum()
+        loss.backward()
+        grad = x.grad
+        assert grad is not None
+        # gradient should be finite everywhere
+        assert torch.all(torch.isfinite(grad))
+        # gradient should be non-zero for non-zero inputs
+        nonzero_mask = self.x_np != 0.0
+        assert torch.all(grad[nonzero_mask] != 0.0)
+
+    def test_silut_make_fx(self) -> None:
+        """make_fx can trace through silut activation."""
+
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            return _torch_activation(x, "silut:10.0")
+
+        traced = make_fx(fn)(self.x_torch)
+        result = traced(self.x_torch)
+        expected = _torch_activation(self.x_torch, "silut:10.0")
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected.detach().numpy(), rtol=1e-12, atol=1e-12
+        )
+
+    def test_silut_below_threshold_is_silu(self) -> None:
+        """Below threshold, silut equals silu exactly."""
+        threshold = 10.0
+        x_below = torch.tensor([-5.0, 0.0, 1.0, 5.0, 9.9], dtype=torch.float64)
+        result = _torch_activation(x_below, "silut:10.0")
+        silu = x_below * torch.sigmoid(x_below)
+        np.testing.assert_allclose(
+            result.detach().numpy(), silu.detach().numpy(), rtol=1e-14, atol=1e-14
+        )
+
+    def test_silut_above_threshold_is_tanh_branch(self) -> None:
+        """Above threshold, silut equals tanh(slope*(x-T))+const."""
+        import math
+
+        threshold = 3.0
+        sig_t = 1.0 / (1.0 + math.exp(-threshold))
+        slope = sig_t + threshold * sig_t * (1.0 - sig_t)
+        const = threshold * sig_t
+
+        x_above = torch.tensor([3.5, 5.0, 10.0, 20.0], dtype=torch.float64)
+        result = _torch_activation(x_above, "silut:3.0")
+        expected = torch.tanh(slope * (x_above - threshold)) + const
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected.detach().numpy(), rtol=1e-14, atol=1e-14
+        )
+
+    def test_silut_export(self) -> None:
+        """torch.export.export can trace through silut activation."""
+
+        class SilutModule(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return _torch_activation(x, "silut:10.0")
+
+        mod = SilutModule()
+        exported = torch.export.export(mod, (self.x_torch,))
+        result = exported.module()(self.x_torch)
+        expected = _torch_activation(self.x_torch, "silut:10.0")
+        np.testing.assert_allclose(
+            result.detach().numpy(), expected.detach().numpy(), rtol=1e-12, atol=1e-12
+        )