From 79598241f997acba405f20bca5f32a4d3b02997f Mon Sep 17 00:00:00 2001
From: luozijian <luozijian0924@gmail.com>
Date: Sat, 5 Sep 2026 11:41:49 +0800
Subject: [PATCH] fix(dpmodel): handle zero-node DPA4C spin graphs

---
 deepmd/dpmodel/descriptor/dpa4c_nn/spin.py    |  5 ++-
 .../common/dpmodel/test_descriptor_dpa4c.py   | 26 +++++++++++++
 source/tests/pt_expt/descriptor/test_dpa4c.py | 39 +++++++++++++++++++
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py
index 56ea73389f..9bb714af6f 100644
--- a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py
+++ b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py
@@ -327,7 +327,10 @@ def call(self, spin_moments: Array, degree_two: Array) -> Array:
                 # anisotropy sum over neighbours.
                 xp.reshape(
                     xp.matmul(xp.permute_dims(quadrupole, (0, 2, 1)), degree_two),
-                    (spin_moments.shape[0], -1),
+                    (
+                        spin_moments.shape[0],
+                        self.quadrupole_width * self.degree_channels[2],
+                    ),
                 ),
                 magnitude,
                 coordination,
diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py
index 4c99eb902d..abe2b63b97 100644
--- a/source/tests/common/dpmodel/test_descriptor_dpa4c.py
+++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py
@@ -804,6 +804,32 @@ def make_spin_descriptor(**overrides: Any) -> DescrptDPA4C:
     return descriptor
 
 
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+@pytest.mark.parametrize("channels", [8, 32])
+@pytest.mark.parametrize("node_count", [0, 2])
+def test_native_spin_graph_without_edges(
+    precision: str, channels: int, node_count: int
+) -> None:
+    """Both zero-node graphs and isolated atoms retain the feature width."""
+    descriptor = make_spin_descriptor(precision=precision, channels=channels)
+    graph = neighbor_graph.NeighborGraph(
+        n_node=np.array([node_count], dtype=np.int64),
+        n_local=np.array([node_count], dtype=np.int64),
+        edge_index=np.empty((2, 0), dtype=np.int64),
+        edge_vec=np.empty((0, 3), dtype=precision),
+        edge_mask=np.empty((0,), dtype=bool),
+    )
+    output, rot_mat = descriptor.call_graph(
+        graph,
+        np.arange(node_count, dtype=np.int64) % 2,
+        spin=np.ones((node_count, 3), dtype=precision),
+    )
+    assert output.shape == (node_count, descriptor.get_dim_out())
+    assert output.dtype == graph.edge_vec.dtype
+    assert np.isfinite(output).all()
+    assert rot_mat is None
+
+
 def spin_reference_terms(
     descriptor: DescrptDPA4C,
     graph: Any,
diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py
index 12ef367c70..df89986821 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4c.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4c.py
@@ -458,6 +458,45 @@ def test_coincident_edge_has_finite_third_derivative(self) -> None:
             )
 
 
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+@pytest.mark.parametrize("channels", [8, 32])
+@pytest.mark.parametrize("node_count", [0, 2])
+def test_native_spin_graph_without_edges(
+    precision: str, channels: int, node_count: int
+) -> None:
+    """The reference spin path supports zero nodes as well as isolated atoms."""
+    descriptor = TestDPA4C.build(
+        precision=precision, channels=channels, use_spin=[True, False]
+    ).eval()
+    with torch.no_grad():
+        descriptor.spin.spin_gate.fill_(1.0)
+    dtype = getattr(torch, precision)
+    graph = NeighborGraph(
+        n_node=torch.tensor([node_count], dtype=torch.int64, device=env.DEVICE),
+        n_local=torch.tensor([node_count], dtype=torch.int64, device=env.DEVICE),
+        edge_index=torch.empty((2, 0), dtype=torch.int64, device=env.DEVICE),
+        edge_vec=torch.empty((0, 3), dtype=dtype, device=env.DEVICE),
+        edge_mask=torch.empty((0,), dtype=torch.bool, device=env.DEVICE),
+    )
+    spin = torch.ones(
+        (node_count, 3), dtype=dtype, device=env.DEVICE, requires_grad=True
+    )
+    output, rot_mat = descriptor.call_graph(
+        graph,
+        torch.arange(node_count, dtype=torch.int64, device=env.DEVICE) % 2,
+        spin=spin,
+    )
+    assert output.shape == (node_count, descriptor.get_dim_out())
+    assert output.dtype == dtype
+    assert output.device == spin.device
+    assert torch.isfinite(output).all()
+    assert rot_mat is None
+    output.sum().backward()
+    assert spin.grad is not None
+    assert spin.grad.shape == spin.shape
+    assert torch.isfinite(spin.grad).all()
+
+
 class TestDPA4CSpinGate:
     """Torch-side contracts of the spin branch gate."""