feat(dpa4): so3_readout across pt + dpmodel/pt_expt backends

deepmd/dpmodel/descriptor/dpa4.py

@@ -194,6 +194,7 @@ def __init__(
         node_wise_so3: bool = False,
         message_node_s2: bool = False,
         message_node_so3: bool = False,
+        so3_readout: str = "none",
         lebedev_quadrature: bool | list[bool] | None = True,
         activation_function: str = "silu",
         glu_activation: bool = True,
@@ -275,6 +276,9 @@ def __init__(
         self.node_wise_so3 = bool(node_wise_so3)
         self.message_node_s2 = bool(message_node_s2)
         self.message_node_so3 = bool(message_node_so3)
+        self.so3_readout = str(so3_readout).lower()
+        if self.so3_readout not in {"none", "glu", "mlp"}:
+            raise ValueError("`so3_readout` must be one of 'none', 'glu', or 'mlp'")
         if lebedev_quadrature is None:
             lebedev_quadrature = [True, True]
         elif isinstance(lebedev_quadrature, bool):
@@ -639,12 +643,20 @@ def __init__(
         self.blocks = blocks
 
         # === Final FFN for l=0 output mixing (fp32+) ===
+        # ``so3_readout="none"`` runs a degree-0 scalar FFN on the l=0 slice.
+        # ``"glu"``/``"mlp"`` run a full FFN at the last block's node degree whose
+        # SO(3) Wigner-D grid folds l>0 geometry into l=0; the value selects the
+        # quadratic grid product or the point-wise grid MLP.
+        readout_lmax = self.node_l_schedule[-1]
         self.output_ffn = EquivariantFFN(
-            lmax=0,
+            lmax=0 if self.so3_readout == "none" else readout_lmax,
             channels=self.channels,
             hidden_channels=self.out_ffn_neurons,
-            grid_mlp=False,
+            kmax=min(self.kmax, readout_lmax),
+            grid_mlp=self.so3_readout == "mlp",
+            grid_branch=0,
             s2_activation=False,
+            ffn_so3_grid=self.so3_readout != "none",
             activation_function=self.out_activation_function,
             glu_activation=self.out_glu_activation,
             mlp_bias=self.mlp_bias,
@@ -926,10 +938,24 @@ def call(
             x = block(x, edge_cache, rad_feat_per_block[i])[0]
 
         # === Step 10. Final l=0 output mixing ===
-        x_scalar = xp.reshape(
-            x[:, 0:1, :, :], (n_nodes, 1, 1, self.channels)
-        )  # (N, 1, 1, C)
-        x_scalar = x_scalar + self.output_ffn(x_scalar)
+        # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full
+        # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The
+        # residual is added on the full coefficient tensor before extracting
+        # l=0 to mirror pt.
+        compute_prec = get_xp_precision(xp, self.compute_precision)
+        if self.so3_readout == "none":
+            ffn_in = xp.astype(
+                xp.reshape(x[:, 0:1, :, :], (n_nodes, 1, 1, self.channels)),
+                compute_prec,
+            )  # (N, 1, 1, C)
+        else:
+            # truncate to the final node degree (what output_ffn is built for);
+            # no-op in the normal path (blocks already shrank x), defensive vs
+            # any path that leaves x at the initial degree. Mirrors pt.
+            ffn_in = xp.astype(
+                x[:, : self.node_ebed_dims[-1], :, :], compute_prec
+            )  # (N, D, 1, C)
+        x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :]
 
         # === Step 11. Reshape and return ===
         descriptor = xp.reshape(x_scalar, (nf, nloc, self.channels))
@@ -1231,6 +1257,7 @@ def serialize(self) -> dict[str, Any]:
                 "node_wise_so3": self.node_wise_so3,
                 "message_node_s2": self.message_node_s2,
                 "message_node_so3": self.message_node_so3,
+                "so3_readout": self.so3_readout,
                 "lebedev_quadrature": self.lebedev_quadrature,
                 "activation_function": self.activation_function,
                 "glu_activation": self.glu_activation,

deepmd/pt/model/descriptor/sezm.py

@@ -320,6 +320,16 @@ class DescrptSeZM(BaseDescriptor, nn.Module):
     message_node_so3
         If True, use the corresponding post-aggregation SO(3) Wigner-D grid-net
         branch. The message is the query and the node state is the context.
+    so3_readout
+        Read-out FFN mode for the final ``l=0`` descriptor. ``"none"`` applies a
+        degree-0 scalar FFN to the ``l=0`` slice only; ``l>0`` coefficients are
+        discarded before the read-out. ``"glu"`` and ``"mlp"`` apply a full
+        equivariant FFN whose degree equals the node degree of the last
+        interaction block, driven by the SO(3) Wigner-D grid, so ``l>0`` geometry
+        is folded into ``l=0`` before the scalar is extracted. The value selects
+        the quadratic grid product (``"glu"``) or the polynomial point-wise grid
+        MLP (``"mlp"``). The Wigner-D frame order follows ``kmax``. The residual
+        stays on the ``l=0`` channel.
     lebedev_quadrature
         Either one boolean applied to both S2 branches, or two booleans
         ``[so2_enabled, ffn_enabled]`` aligned with ``s2_activation``. If
@@ -425,6 +435,7 @@ def __init__(
         node_wise_so3: bool = False,
         message_node_s2: bool = False,
         message_node_so3: bool = False,
+        so3_readout: str = "none",
         lebedev_quadrature: bool | list[bool] | None = True,
         activation_function: str = "silu",
         glu_activation: bool = True,
@@ -512,6 +523,9 @@ def __init__(
         self.node_wise_so3 = bool(node_wise_so3)
         self.message_node_s2 = bool(message_node_s2)
         self.message_node_so3 = bool(message_node_so3)
+        self.so3_readout = str(so3_readout).lower()
+        if self.so3_readout not in {"none", "glu", "mlp"}:
+            raise ValueError("`so3_readout` must be one of 'none', 'glu', or 'mlp'")
         if lebedev_quadrature is None:
             lebedev_quadrature = [True, True]
         elif isinstance(lebedev_quadrature, bool):
@@ -932,13 +946,21 @@ def __init__(
             )
 
         # === Final FFN for l=0 output mixing ===
+        # ``so3_readout="none"`` runs a degree-0 scalar FFN on the l=0 slice.
+        # ``"glu"``/``"mlp"`` run a full FFN at the last block's node degree whose
+        # SO(3) Wigner-D grid folds l>0 geometry into l=0; the value selects the
+        # quadratic grid product or the point-wise grid MLP.
+        readout_lmax = self.node_l_schedule[-1]
         self.output_ffn = EquivariantFFN(
-            lmax=0,
+            lmax=0 if self.so3_readout == "none" else readout_lmax,
             channels=self.channels,
             hidden_channels=self.out_ffn_neurons,
-            grid_mlp=False,
+            kmax=min(self.kmax, readout_lmax),
+            grid_mlp=self.so3_readout == "mlp",
+            grid_branch=0,
             dtype=self.compute_dtype,
             s2_activation=False,
+            ffn_so3_grid=self.so3_readout != "none",
             activation_function=self.out_activation_function,
             glu_activation=self.out_glu_activation,
             mlp_bias=self.mlp_bias,
@@ -1205,15 +1227,23 @@ def forward(
                     x = self._forward_blocks(x, edge_cache, rad_feat_per_block)
 
         # === Step 11. Final l=0 output mixing ===
-        # Extract l=0 scalar features and apply FFN in promoted dtype.
-        # Residual keeps the output close to identity with zero-initialized FFN output.
+        # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full
+        # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The
+        # residual is added on the full coefficient tensor before extracting
+        # l=0: slicing the summed tensor rather than the FFN output keeps the
+        # saved degree-axis stride static under torch.compile dynamic shapes.
         with nvtx_range("output_ffn"):
-            x_scalar = (
+            ffn_in = (
                 x[:, 0:1, :, :]
                 .reshape(n_nodes, 1, 1, self.channels)
                 .to(dtype=self.compute_dtype)
-            )  # (N, 1, 1, C)
-            x_scalar = x_scalar + self.output_ffn(x_scalar)
+                if self.so3_readout == "none"
+                # truncate to the final node degree: the empty-edge path
+                # skips the blocks, leaving x at node_ebed_dims[0]; output_ffn
+                # is built for node_ebed_dims[-1]. No-op when blocks ran.
+                else x[:, : self.node_ebed_dims[-1], :, :].to(dtype=self.compute_dtype)
+            )
+            x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :]
 
         # === Step 12. Reshape to (nf, nloc, channels) and return ===
         descriptor = rearrange(
@@ -1380,13 +1410,23 @@ def forward_with_edges(
                 x = self._forward_blocks(x, edge_cache, rad_feat_per_block)
 
         # === Step 10. Final l=0 output mixing ===
+        # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full
+        # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The
+        # residual is added on the full coefficient tensor before extracting
+        # l=0: slicing the summed tensor rather than the FFN output keeps the
+        # saved degree-axis stride static under torch.compile dynamic shapes.
         with nvtx_range("output_ffn"):
-            x_scalar = (
+            ffn_in = (
                 x[:, 0:1, :, :]
                 .reshape(n_nodes, 1, 1, self.channels)
                 .to(dtype=self.compute_dtype)
-            )  # (N, 1, 1, C)
-            x_scalar = x_scalar + self.output_ffn(x_scalar)
+                if self.so3_readout == "none"
+                # truncate to the final node degree: the empty-edge path
+                # skips the blocks, leaving x at node_ebed_dims[0]; output_ffn
+                # is built for node_ebed_dims[-1]. No-op when blocks ran.
+                else x[:, : self.node_ebed_dims[-1], :, :].to(dtype=self.compute_dtype)
+            )
+            x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :]
 
         # === Step 11. Reshape to (nf, nloc, channels) and return ===
         descriptor = x_scalar.reshape(nf, nloc, self.channels)  # (nf, nloc, C)
@@ -2043,6 +2083,7 @@ def serialize(self) -> dict[str, Any]:
                 "node_wise_so3": self.node_wise_so3,
                 "message_node_s2": self.message_node_s2,
                 "message_node_so3": self.message_node_so3,
+                "so3_readout": self.so3_readout,
                 "lebedev_quadrature": self.lebedev_quadrature,
                 "activation_function": self.activation_function,
                 "glu_activation": self.glu_activation,

deepmd/utils/argcheck.py

@@ -551,6 +551,16 @@ def descrpt_se_zm_args() -> list[Argument]:
         "context. When enabled together with `message_node_s2`, the SO(3) "
         "branch is used for this path."
     )
+    doc_so3_readout = (
+        "Read-out FFN mode for the final l=0 descriptor. `none` applies a "
+        "degree-0 scalar FFN to the l=0 slice only; l>0 coefficients are "
+        "discarded before the read-out. `glu` and `mlp` apply a full equivariant "
+        "FFN on the SO(3) Wigner-D grid so l>0 geometry is folded into l=0 "
+        "before the scalar is extracted; the value selects the quadratic grid "
+        "product (`glu`) or the polynomial point-wise grid MLP (`mlp`). The "
+        "read-out degree equals the node degree of the last interaction block; "
+        "the Wigner-D frame order follows `kmax`."
+    )
     doc_lebedev_quadrature = (
         "Either one boolean applied to both S2 branches, or two booleans "
         "`[so2_enabled, ffn_enabled]` aligned with `s2_activation`. If a branch "
@@ -881,6 +891,15 @@ def descrpt_se_zm_args() -> list[Argument]:
             default=False,
             doc=doc_only_pt_supported + doc_message_node_so3,
         ),
+        Argument(
+            "so3_readout",
+            str,
+            optional=True,
+            default="none",
+            extra_check=lambda x: x in ("none", "glu", "mlp"),
+            extra_check_errmsg="must be one of 'none', 'glu', or 'mlp'",
+            doc=doc_only_pt_supported + doc_so3_readout,
+        ),
         Argument(
             "lebedev_quadrature",
             [bool, list[bool]],

examples/water/dpa4/input.json

@@ -23,6 +23,7 @@
       "n_atten_head": 1,
       "ffn_neurons": 0,
       "ffn_so3_grid": true,
+      "so3_readout": "mlp",
       "grid_mlp": [
         false,
         false,

source/tests/consistent/descriptor/test_dpa4.py

@@ -49,6 +49,7 @@
     "ffn_so3_grid",
     "message_node_so3",
     "grid_mlp",
+    "so3_readout",
 )
 
 DPA4_BASELINE_CASE = {
@@ -59,6 +60,7 @@
     "ffn_so3_grid": False,
     "message_node_so3": False,
     "grid_mlp": False,
+    "so3_readout": "none",
 }
 
 
@@ -99,6 +101,10 @@ def dpa4_case(**overrides: Any) -> tuple:
     dpa4_case(ffn_so3_grid=True, message_node_so3=True),
     # polynomial grid MLP op (grid_branch=0 so grid_mlp takes effect)
     dpa4_case(grid_mlp=True, grid_branch=[0, 0, 0]),
+    # SO(3) readout: GLU grid product folds l>0 into the l=0 output
+    dpa4_case(so3_readout="glu"),
+    # SO(3) readout: point-wise grid MLP folds l>0 into the l=0 output
+    dpa4_case(so3_readout="mlp"),
 )
 
 
@@ -114,6 +120,7 @@ def data(self) -> dict:
             ffn_so3_grid,
             message_node_so3,
             grid_mlp,
+            so3_readout,
         ) = self.param
         return {
             "ntypes": self.ntypes,
@@ -130,6 +137,7 @@ def data(self) -> dict:
             "ffn_so3_grid": ffn_so3_grid,
             "message_node_so3": message_node_so3,
             "grid_mlp": grid_mlp,
+            "so3_readout": so3_readout,
             "random_gamma": False,
             "precision": precision,
             "trainable": False,

source/tests/pt/model/test_descriptor_sezm.py

@@ -155,6 +155,30 @@ def _assert_forward_backward_smoke(self, **model_kwargs) -> DescrptSeZM:
         self.assertTrue(torch.all(torch.isfinite(extended_coord.grad)))
         return model
 
+    def test_so3_readout_empty_edge_shrinking_schedule(self) -> None:
+        """so3_readout glu/mlp must handle the empty-edge path.
+
+        With a shrinking ``l_schedule`` and no edges (every atom isolated),
+        ``_forward_blocks`` is skipped so ``x`` keeps the *initial* node degree
+        ``node_ebed_dims[0]``; the readout must truncate it to the final degree
+        ``node_ebed_dims[-1]`` (what ``output_ffn`` is built for) before the FFN.
+        Regression for the readout shape mismatch on isolated atoms.
+        """
+        coord, atype, _ = _tiny_two_atom_system(self.device, dtype=torch.float32)
+        extended_coord = coord.reshape(1, -1).detach().requires_grad_(True)
+        # all neighbors masked out -> edge_cache.src.numel() == 0 -> blocks skipped
+        nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device)
+        for readout in ("glu", "mlp"):
+            with self.subTest(so3_readout=readout):
+                model = DescrptSeZM(
+                    **_descriptor_kwargs(l_schedule=[2, 1], so3_readout=readout)
+                )
+                desc, *_ = model(
+                    extended_coord, atype, nlist, mapping=None, comm_dict=None
+                )
+                self.assertEqual(desc.shape, (1, 2, 4))
+                self.assertTrue(torch.all(torch.isfinite(desc)))
+
     def test_forward_with_descriptor_variants(self) -> None:
         """Test forward/backward smoke paths for compact descriptor variants."""
         cases = {

source/tests/pt/model/test_dpa4_dpmodel_parity.py

@@ -3578,6 +3578,60 @@ def test_descriptor_extra_node_l(self) -> None:
         pt_mod, dp_mod, _ = self._build_descr_pair(extra_node_l=1)
         self._assert_descr_parity(pt_mod, dp_mod)
 
+    @pytest.mark.parametrize(
+        "so3_readout", ["glu", "mlp"]
+    )  # SO(3) grid readout: quadratic grid product vs point-wise grid MLP
+    def test_descriptor_so3_readout(self, so3_readout) -> None:
+        # so3_readout!="none" feeds the full (N, D, 1, C) node tensor to the
+        # output FFN so the SO(3) Wigner-D grid folds l>0 into l=0. The pt
+        # reference is pinned to CPU so the parity holds under the CUDA default
+        # device; the gate stays at the strict fp64 descriptor tolerance.
+        from deepmd.dpmodel.descriptor.dpa4 import (
+            DescrptDPA4,
+        )
+        from deepmd.pt.model.descriptor.sezm import (
+            DescrptSeZM,
+        )
+
+        kwargs = self._descr_kwargs(so3_readout=so3_readout)
+        pt_mod = DescrptSeZM(**kwargs).double().eval().to("cpu")
+        # so3_linear_2 / output projections are zero-initialized; perturb so the
+        # readout output is nontrivial (otherwise it is identically ~0)
+        rng = np.random.default_rng(2160)
+        with torch.no_grad():
+            for p in pt_mod.parameters():
+                p += torch.from_numpy(0.05 * rng.normal(size=tuple(p.shape))).to("cpu")
+        dp_mod = DescrptDPA4.deserialize(pt_mod.serialize())
+        assert dp_mod.so3_readout == so3_readout
+
+        inp = self._inputs()
+        coord, atype_ext, nlist, mp = (
+            inp["coord"],
+            inp["atype_ext"],
+            inp["nlist"],
+            inp["mapping"],
+        )
+        nf = coord.shape[0]
+        out_dp = np.asarray(
+            dp_mod.call(coord.reshape(nf, -1), atype_ext, nlist, mapping=mp)[0]
+        )
+        out_pt = (
+            pt_mod(
+                torch.from_numpy(coord).to("cpu"),
+                torch.from_numpy(atype_ext.astype(np.int64)).to("cpu"),
+                torch.from_numpy(nlist.astype(np.int64)).to("cpu"),
+                mapping=torch.from_numpy(mp.astype(np.int64)).to("cpu"),
+            )[0]
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        assert out_dp.shape == out_pt.shape
+        # nontrivial output magnitude (guards against a trivially-zero readout)
+        assert np.abs(out_dp).max() > 1e-6
+        # strict fp64 descriptor-level gate
+        np.testing.assert_allclose(out_dp, out_pt, rtol=1e-10, atol=1e-12)
+
     def test_descriptor_torch_namespace(self) -> None:
         # the dp descriptor must run under the torch array namespace as well:
         # feeding torch tensors must yield a torch tensor matching the numpy