From e334724c348fd13564b02280acf5ff30beb1d885 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sat, 10 Oct 2026 15:44:26 +0800
Subject: [PATCH 1/7] perf(pt): assemble Wigner-D behind a custom op under
 compile

Functionalization expands the in-place block scatter into a chain of
full-size zeros/index_put adds plus a realized transpose (~4 GB of HBM
traffic per frame at the SeZM N=4096 workload).  A wigner_d_assemble
custom op keeps the eager scatter and hands the graph two contiguous
tensors; the traced assembly takes it whenever a fake-tensor mode is
active, so the eager path is untouched.
---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 163 +++++++++++++++---
 1 file changed, 139 insertions(+), 24 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index b283e3a7e5..d420dec53d 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -45,6 +45,88 @@
 _TORCH_RELEASE = Version(torch.__version__).release[:2]
 
 
+def _wigner_d_assemble_impl(
+    blocks: list[torch.Tensor],
+    dim_full: int,
+    poly_lmin: int,
+    n_regular: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Assemble the dense block-diagonal matrix and its contiguous transpose.
+
+    In-place slice assignments are exactly what Inductor's functionalization
+    cannot keep in place: traced, they expand into a chain of full-size
+    ``zeros + index_put + add`` tensors plus a realized transpose (~4 GB of
+    HBM traffic at the SeZM N=4096 workload).  Behind a custom op the eager
+    scatter stays a scatter -- one buffer, one transpose copy.
+    """
+    ref = blocks[0]
+    n_edge = ref.shape[0]
+    d_full = torch.zeros(n_edge, dim_full, dim_full, dtype=ref.dtype, device=ref.device)
+    d_full[:, 0, 0] = 1.0
+    off = 1
+    for block in blocks[:n_regular]:
+        size = block.shape[1]
+        d_full[:, off : off + size, off : off + size] = block
+        off += size
+    if n_regular < len(blocks):
+        po = poly_lmin * poly_lmin
+        d_full[:, po:, po:] = blocks[n_regular]
+    return d_full, d_full.transpose(-1, -2).contiguous()
+
+
+_wigner_d_assemble_op = torch.library.custom_op(
+    "sezm_nn::wigner_d_assemble", mutates_args=()
+)(_wigner_d_assemble_impl)
+
+
+@_wigner_d_assemble_op.register_fake
+def _(
+    blocks: list[torch.Tensor],
+    dim_full: int,
+    poly_lmin: int,
+    n_regular: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ref = blocks[0]
+    d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full))
+    return d_full, d_full.transpose(-1, -2).contiguous()
+
+
+def _wigner_d_assemble_setup(
+    ctx: Any, inputs: Any, output: Any
+) -> None:
+    blocks, _, poly_lmin, n_regular = inputs
+    ctx.blocks = blocks
+    ctx.poly_lmin = poly_lmin
+    ctx.n_regular = n_regular
+
+
+def _wigner_d_assemble_backward(
+    ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor
+) -> tuple[list[torch.Tensor], None, None, None]:
+    # Each block lands in D as itself and in Dt transposed, so its gradient
+    # is the matching D-slice plus the transpose of the matching Dt-slice.
+    # The [0, 0] constant receives no gradient.
+    grads = []
+    off = 1
+    for block in ctx.blocks[: ctx.n_regular]:
+        size = block.shape[1]
+        g = grad_d_full[:, off : off + size, off : off + size]
+        gt = grad_dt_full[:, off : off + size, off : off + size]
+        grads.append(g + gt.transpose(-1, -2))
+        off += size
+    if ctx.n_regular < len(ctx.blocks):
+        po = ctx.poly_lmin * ctx.poly_lmin
+        g = grad_d_full[:, po:, po:]
+        gt = grad_dt_full[:, po:, po:]
+        grads.append(g + gt.transpose(-1, -2))
+    return grads, None, None, None
+
+
+_wigner_d_assemble_op.register_autograd(
+    _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup
+)
+
+
 class CaseCoefficients(nn.Module):
     """
     Polynomial tables for one magnitude-ordered branch of the quaternion Wigner path.
@@ -546,62 +628,63 @@ def forward(
             eps=self.eps,
         )
         n_edge = edge_quaternion.shape[0]
-        D_full = torch.zeros(
-            n_edge,
-            self.dim_full,
-            self.dim_full,
-            dtype=edge_quaternion.dtype,
-            device=edge_quaternion.device,
-        )
-        D_full[:, 0, 0] = 1.0
+
+        # Degree blocks are computed first; the dense assembly then either
+        # runs in place (eager) or behind the opaque wigner_d_assemble op
+        # (compiled).  Under tracing, the in-place slice assignments below
+        # would be functionalized into a chain of full-size zeros/index_put
+        # adds plus a realized transpose; the custom op keeps the eager
+        # scatter and hands the graph two plain contiguous tensors.
+        blocks: list[torch.Tensor] = []
+        poly_block: torch.Tensor | None = None
 
         if self.lmax >= 1:
             with nvtx_range("WignerD/l1"):
-                D_full[:, 1:4, 1:4] = self._compute_l1_block(edge_quaternion)
+                blocks.append(self._compute_l1_block(edge_quaternion))
 
         if self.lmax >= 2:
             with nvtx_range("WignerD/l2"):
-                D_full[:, 4:9, 4:9] = self._compute_l2_block(edge_quaternion)
+                blocks.append(self._compute_l2_block(edge_quaternion))
 
         if self.lmax >= 3:
             if self.lmax >= 4:
                 with nvtx_range("WignerD/l3l4"):
                     D_l3, D_l4 = self._compute_l3l4_blocks(edge_quaternion)
-                    D_full[:, 9:16, 9:16] = D_l3
-                    D_full[:, 16:25, 16:25] = D_l4
+                    blocks.append(D_l3)
+                    blocks.append(D_l4)
             else:
                 with nvtx_range("WignerD/l3"):
-                    D_full[:, 9:16, 9:16] = self._compute_l3_block(edge_quaternion)
+                    blocks.append(self._compute_l3_block(edge_quaternion))
 
         if self.lmax >= 5:
             if self.lmax >= 6:
                 with nvtx_range("WignerD/l5l6"):
                     D_l5, D_l6 = self._compute_l5l6_blocks(edge_quaternion)
-                    D_full[:, 25:36, 25:36] = D_l5
-                    D_full[:, 36:49, 36:49] = D_l6
+                    blocks.append(D_l5)
+                    blocks.append(D_l6)
             else:
                 with nvtx_range("WignerD/l5"):
-                    D_full[:, 25:36, 25:36] = self._compute_l5_block(edge_quaternion)
+                    blocks.append(self._compute_l5_block(edge_quaternion))
 
         if self.lmax >= 7:
             if self.lmax >= 8:
                 with nvtx_range("WignerD/l7l8"):
                     D_l7, D_l8 = self._compute_l7l8_blocks(edge_quaternion)
-                    D_full[:, 49:64, 49:64] = D_l7
-                    D_full[:, 64:81, 64:81] = D_l8
+                    blocks.append(D_l7)
+                    blocks.append(D_l8)
             else:
                 with nvtx_range("WignerD/l7"):
-                    D_full[:, 49:64, 49:64] = self._compute_l7_block(edge_quaternion)
+                    blocks.append(self._compute_l7_block(edge_quaternion))
 
         if self.lmax >= 9:
             if self.lmax >= 10:
                 with nvtx_range("WignerD/l9l10"):
                     D_l9, D_l10 = self._compute_l9l10_blocks(edge_quaternion)
-                    D_full[:, 81:100, 81:100] = D_l9
-                    D_full[:, 100:121, 100:121] = D_l10
+                    blocks.append(D_l9)
+                    blocks.append(D_l10)
             else:
                 with nvtx_range("WignerD/l9"):
-                    D_full[:, 81:100, 81:100] = self._compute_l9_block(edge_quaternion)
+                    blocks.append(self._compute_l9_block(edge_quaternion))
 
         if self.lmax >= self.poly_lmin:
             with nvtx_range("WignerD/polynomial"):
@@ -616,7 +699,7 @@ def forward(
                     self.poly_coeffs,
                     dtype=self.dtype,
                 )
-                D_poly = self._wigner_d_pair_to_real(
+                poly_block = self._wigner_d_pair_to_real(
                     D_re,
                     D_im,
                     (
@@ -628,7 +711,39 @@ def forward(
                     lmax=self.lmax,
                     lmin=self.poly_lmin,
                 )
-                D_full[:, self.poly_offset :, self.poly_offset :] = D_poly
+
+        # ``is_compiling`` alone is not enough here: this module is traced
+        # through make_fx fake tensors, where the flag stays False.  A fake
+        # tensor carries its mode, so probe for that instead.
+        tracing = (
+            torch.compiler.is_compiling()
+            or getattr(blocks[0], "fake_mode", None) is not None
+            if blocks
+            else False
+        )
+        if blocks and tracing:
+            return _wigner_d_assemble_op(
+                blocks + ([poly_block] if poly_block is not None else []),
+                self.dim_full,
+                self.poly_lmin,
+                len(blocks),
+            )
+
+        D_full = torch.zeros(
+            n_edge,
+            self.dim_full,
+            self.dim_full,
+            dtype=edge_quaternion.dtype,
+            device=edge_quaternion.device,
+        )
+        D_full[:, 0, 0] = 1.0
+        off = 1
+        for block in blocks:
+            size = block.shape[1]
+            D_full[:, off : off + size, off : off + size] = block
+            off += size
+        if poly_block is not None:
+            D_full[:, self.poly_offset :, self.poly_offset :] = poly_block
 
         # Consumers address the inverse rotation through explicit strides or
         # PyTorch strided operators, so the transpose can share D_full's storage.

From aea99479029413e4c53ddeed59d0327cec86d763 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sat, 10 Oct 2026 15:58:33 +0800
Subject: [PATCH 2/7] perf(pt): return the inverse Wigner-D as a view, not a
 copy

The flash and rotation-back kernels already address the rotation
through explicit strides, so the contiguous transpose the assembly
produced was pure storage: one 660 MB copy per frame that existed only
to be read back.  The assemble op now returns D_full alone and the
caller hands out a transpose view; consumers read it in place and no
transpose copy ever exists.
---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 42 +++++++++----------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index d420dec53d..1c5e3d9f8c 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -50,14 +50,18 @@ def _wigner_d_assemble_impl(
     dim_full: int,
     poly_lmin: int,
     n_regular: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Assemble the dense block-diagonal matrix and its contiguous transpose.
+) -> torch.Tensor:
+    """Assemble the dense block-diagonal matrix; the transpose stays a view.
 
     In-place slice assignments are exactly what Inductor's functionalization
     cannot keep in place: traced, they expand into a chain of full-size
     ``zeros + index_put + add`` tensors plus a realized transpose (~4 GB of
     HBM traffic at the SeZM N=4096 workload).  Behind a custom op the eager
-    scatter stays a scatter -- one buffer, one transpose copy.
+    scatter stays a scatter.  Only ``D_full`` is produced here; the caller
+    makes ``Dt_full`` a plain transpose view outside the op, and every
+    consumer addresses it through strides (the flash and rotation kernels
+    take explicit stride arguments; the eager consumers are strided torch
+    operators).
     """
     ref = blocks[0]
     n_edge = ref.shape[0]
@@ -71,7 +75,7 @@ def _wigner_d_assemble_impl(
     if n_regular < len(blocks):
         po = poly_lmin * poly_lmin
         d_full[:, po:, po:] = blocks[n_regular]
-    return d_full, d_full.transpose(-1, -2).contiguous()
+    return d_full
 
 
 _wigner_d_assemble_op = torch.library.custom_op(
@@ -85,15 +89,12 @@ def _(
     dim_full: int,
     poly_lmin: int,
     n_regular: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> torch.Tensor:
     ref = blocks[0]
-    d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full))
-    return d_full, d_full.transpose(-1, -2).contiguous()
+    return ref.new_zeros((ref.shape[0], dim_full, dim_full))
 
 
-def _wigner_d_assemble_setup(
-    ctx: Any, inputs: Any, output: Any
-) -> None:
+def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None:
     blocks, _, poly_lmin, n_regular = inputs
     ctx.blocks = blocks
     ctx.poly_lmin = poly_lmin
@@ -101,24 +102,20 @@ def _wigner_d_assemble_setup(
 
 
 def _wigner_d_assemble_backward(
-    ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor
+    ctx: Any, grad_d_full: torch.Tensor
 ) -> tuple[list[torch.Tensor], None, None, None]:
-    # Each block lands in D as itself and in Dt transposed, so its gradient
-    # is the matching D-slice plus the transpose of the matching Dt-slice.
-    # The [0, 0] constant receives no gradient.
+    # The transpose view's gradient reaches D_full through the ordinary
+    # autograd of aten.transpose, so only D-slices arrive here.  The [0, 0]
+    # constant receives no gradient.
     grads = []
     off = 1
     for block in ctx.blocks[: ctx.n_regular]:
         size = block.shape[1]
-        g = grad_d_full[:, off : off + size, off : off + size]
-        gt = grad_dt_full[:, off : off + size, off : off + size]
-        grads.append(g + gt.transpose(-1, -2))
+        grads.append(grad_d_full[:, off : off + size, off : off + size])
         off += size
     if ctx.n_regular < len(ctx.blocks):
         po = ctx.poly_lmin * ctx.poly_lmin
-        g = grad_d_full[:, po:, po:]
-        gt = grad_dt_full[:, po:, po:]
-        grads.append(g + gt.transpose(-1, -2))
+        grads.append(grad_d_full[:, po:, po:])
     return grads, None, None, None
 
 
@@ -722,12 +719,15 @@ def forward(
             else False
         )
         if blocks and tracing:
-            return _wigner_d_assemble_op(
+            d_full = _wigner_d_assemble_op(
                 blocks + ([poly_block] if poly_block is not None else []),
                 self.dim_full,
                 self.poly_lmin,
                 len(blocks),
             )
+            # Stride-addressed consumers read the inverse rotation straight
+            # from the shared storage; no transpose copy ever exists.
+            return d_full, d_full.transpose(-1, -2)
 
         D_full = torch.zeros(
             n_edge,

From 40fac5908d005fda5fbb90e297e30c170a25dc4d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 15 Oct 2026 20:34:22 +0800
Subject: [PATCH 3/7] refactor(pt): share the Wigner-D assembly and
 inverse-layout exit, drop tensor retention

---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 86 ++++++++++---------
 1 file changed, 46 insertions(+), 40 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index 1c5e3d9f8c..f9be1ab451 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -96,9 +96,12 @@ def _(
 
 def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None:
     blocks, _, poly_lmin, n_regular = inputs
-    ctx.blocks = blocks
+    # Only the layout metadata is needed to slice the cotangent; storing the
+    # block tensors would extend every activation's lifetime past the
+    # forward, which the backward never reads.
+    ctx.block_widths = [int(block.shape[1]) for block in blocks[:n_regular]]
+    ctx.has_poly_tail = n_regular < len(blocks)
     ctx.poly_lmin = poly_lmin
-    ctx.n_regular = n_regular
 
 
 def _wigner_d_assemble_backward(
@@ -109,11 +112,10 @@ def _wigner_d_assemble_backward(
     # constant receives no gradient.
     grads = []
     off = 1
-    for block in ctx.blocks[: ctx.n_regular]:
-        size = block.shape[1]
+    for size in ctx.block_widths:
         grads.append(grad_d_full[:, off : off + size, off : off + size])
         off += size
-    if ctx.n_regular < len(ctx.blocks):
+    if ctx.has_poly_tail:
         po = ctx.poly_lmin * ctx.poly_lmin
         grads.append(grad_d_full[:, po:, po:])
     return grads, None, None, None
@@ -626,12 +628,13 @@ def forward(
         )
         n_edge = edge_quaternion.shape[0]
 
-        # Degree blocks are computed first; the dense assembly then either
-        # runs in place (eager) or behind the opaque wigner_d_assemble op
-        # (compiled).  Under tracing, the in-place slice assignments below
-        # would be functionalized into a chain of full-size zeros/index_put
-        # adds plus a realized transpose; the custom op keeps the eager
-        # scatter and hands the graph two plain contiguous tensors.
+        # Degree blocks are computed first; the dense assembly below is one
+        # shared implementation, reached directly in eager execution and
+        # behind the opaque wigner_d_assemble operator under tracing (where
+        # functionalization would otherwise expand the in-place scatter into
+        # a chain of full-size zeros/index_put adds plus a realized
+        # transpose).  The operator hands the graph one contiguous tensor;
+        # the inverse rotation leaves through the shared transpose-view exit.
         blocks: list[torch.Tensor] = []
         poly_block: torch.Tensor | None = None
 
@@ -718,41 +721,44 @@ def forward(
             if blocks
             else False
         )
-        if blocks and tracing:
-            d_full = _wigner_d_assemble_op(
-                blocks + ([poly_block] if poly_block is not None else []),
+        # One assembly implementation serves both routes: the eager call
+        # goes straight to the differentiable scatter (autograd flows through
+        # the slice assignments), the traced call hides it behind the opaque
+        # operator so functionalization cannot expand the scatter.
+        all_blocks = blocks + ([poly_block] if poly_block is not None else [])
+        if not all_blocks:
+            # lmax == 0: the rotation acts only on the scalar channel.
+            D_full = torch.zeros(
+                n_edge,
+                self.dim_full,
+                self.dim_full,
+                dtype=edge_quaternion.dtype,
+                device=edge_quaternion.device,
+            )
+            D_full[:, 0, 0] = 1.0
+        elif tracing:
+            D_full = _wigner_d_assemble_op(
+                all_blocks,
                 self.dim_full,
                 self.poly_lmin,
                 len(blocks),
             )
-            # Stride-addressed consumers read the inverse rotation straight
-            # from the shared storage; no transpose copy ever exists.
-            return d_full, d_full.transpose(-1, -2)
-
-        D_full = torch.zeros(
-            n_edge,
-            self.dim_full,
-            self.dim_full,
-            dtype=edge_quaternion.dtype,
-            device=edge_quaternion.device,
-        )
-        D_full[:, 0, 0] = 1.0
-        off = 1
-        for block in blocks:
-            size = block.shape[1]
-            D_full[:, off : off + size, off : off + size] = block
-            off += size
-        if poly_block is not None:
-            D_full[:, self.poly_offset :, self.poly_offset :] = poly_block
-
-        # Consumers address the inverse rotation through explicit strides or
-        # PyTorch strided operators, so the transpose can share D_full's storage.
+        else:
+            D_full = _wigner_d_assemble_impl(
+                all_blocks,
+                self.dim_full,
+                self.poly_lmin,
+                len(blocks),
+            )
+
+        # One inverse-layout exit shared by both routes: the compatibility
+        # policy below is preserved as-is instead of being bypassed by the
+        # traced branch.  Stride-addressed consumers read the inverse
+        # rotation from the shared storage; PyTorch 2.11 Inductor cannot
+        # safely lower that escaping transpose view, so it is materialized
+        # there exactly as the eager route does.
         Dt_full = D_full.transpose(-1, -2)
         if self._materialize_inverse_rotation:
-            # PyTorch 2.11 Inductor cannot safely lower this escaping transpose
-            # view after the slice assignments that assemble D_full. The
-            # materialized layout keeps the compiled graph semantically
-            # identical; later releases retain the shared-storage view.
             Dt_full = Dt_full.contiguous()
         return D_full, Dt_full
 

From e98c691897acd7bd89004eb5633473847775d39d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 15 Oct 2026 20:34:22 +0800
Subject: [PATCH 4/7] test(pt): focused coverage for the Wigner-D assemble
 operator

---
 .../pt/model/test_descriptor_sezm_triton.py   | 206 ++++++++++++++++++
 1 file changed, 206 insertions(+)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f359f073d6..9094065d64 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2053,3 +2053,209 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestSeZMWignerDAssembleOp(unittest.TestCase):
+    """Focused coverage for the op-bound dense Wigner-D assembly."""
+
+    WIDTHS = (3, 5, 7)  # lmax = 3 ladder
+
+    def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64):
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        return [
+            torch.randn(n_edge, w, w, device="cuda", dtype=dtype, generator=generator)
+            for w in widths
+        ]
+
+    def _reference_assemble(self, blocks, poly, dim_full, poly_lmin):
+        ref = blocks[0]
+        dense = ref.new_zeros(ref.shape[0], dim_full, dim_full)
+        dense[:, 0, 0] = 1.0
+        off = 1
+        for block in blocks:
+            size = block.shape[1]
+            dense[:, off : off + size, off : off + size] = block
+            off += size
+        if poly is not None:
+            po = poly_lmin * poly_lmin
+            dense[:, po:, po:] = poly
+        return dense
+
+    def _loss(self, dense, seed):
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        w1 = torch.randn(
+            dense.shape, device="cuda", dtype=dense.dtype, generator=generator
+        )
+        w2 = torch.randn(
+            dense.shape, device="cuda", dtype=dense.dtype, generator=generator
+        )
+        # cotangent contributions through both D and its transpose
+        return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum()
+
+    @_GPU_KERNELS
+    def test_op_matches_reference_with_grads(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            _wigner_d_assemble_impl,
+            _wigner_d_assemble_op,
+        )
+
+        n_edge = 64
+        for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)):
+            with self.subTest(widths=widths, tail=poly_lmin is not None):
+                blocks = self._make_blocks(widths, n_edge, seed=5)
+                poly = None
+                n_regular = len(blocks)
+                if poly_lmin is not None:
+                    poly_w = 4
+                    poly = self._make_blocks([poly_w], n_edge, seed=7)[0]
+                dim_full = sum(widths) + 1 if poly is None else poly_lmin**2 + poly_w
+                args = (
+                    blocks + ([poly] if poly is not None else []),
+                    dim_full,
+                    poly_lmin or 11,
+                    n_regular,
+                )
+
+                leaves_op = [
+                    b.clone().requires_grad_(True)
+                    for b in (blocks + ([poly] if poly is not None else []))
+                ]
+                dense_op = _wigner_d_assemble_op(leaves_op, args[1], args[2], args[3])
+                leaves_ref = [
+                    b.clone().requires_grad_(True)
+                    for b in (blocks + ([poly] if poly is not None else []))
+                ]
+                dense_ref = _wigner_d_assemble_impl(
+                    leaves_ref, args[1], args[2], args[3]
+                )
+                self.assertTrue(torch.equal(dense_op.detach(), dense_ref.detach()))
+                self.assertTrue(
+                    torch.equal(
+                        dense_op.detach(),
+                        self._reference_assemble(blocks, poly, args[1], args[2]),
+                    )
+                )
+
+                grads_op = torch.autograd.grad(
+                    self._loss(dense_op, 11), leaves_op, allow_unused=True
+                )
+                grads_ref = torch.autograd.grad(
+                    self._loss(dense_ref, 11), leaves_ref, allow_unused=True
+                )
+                for g1, g2 in zip(grads_op, grads_ref, strict=True):
+                    torch.testing.assert_close(g1, g2, atol=0.0, rtol=0.0)
+
+                # fp64 gradcheck on the operator route (small edge count)
+                small = [
+                    b[:3].clone().detach().requires_grad_(True)
+                    for b in (blocks + ([poly] if poly is not None else []))
+                ]
+                self.assertTrue(
+                    torch.autograd.gradcheck(
+                        lambda *bs: _wigner_d_assemble_op(
+                            list(bs), args[1], args[2], args[3]
+                        ),
+                        small,
+                    )
+                )
+
+    @_GPU_KERNELS
+    def test_opcheck_schema_fake_autograd(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            _wigner_d_assemble_op,
+        )
+
+        for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)):
+            with self.subTest(widths=widths, tail=poly_lmin is not None):
+                blocks = self._make_blocks(widths, 8, seed=9)
+                args_tail = []
+                n_regular = len(blocks)
+                if poly_lmin is not None:
+                    args_tail = self._make_blocks([4], 8, seed=13)
+                    blocks = blocks + args_tail
+                dim_full = sum(widths) + 1 if not args_tail else poly_lmin**2 + 4
+                saved = torch.get_default_device()
+                torch.set_default_device(None)
+                try:
+                    with torch.device("cuda"):
+                        torch.library.opcheck(
+                            _wigner_d_assemble_op,
+                            (blocks, dim_full, poly_lmin or 11, n_regular),
+                        )
+                finally:
+                    torch.set_default_device(saved)
+
+    @_GPU_KERNELS
+    def test_traced_matches_eager_and_layout_policy(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        generator = torch.Generator(device="cuda").manual_seed(21)
+        quat = torch.randn(32, 4, device="cuda", generator=generator)
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+        calc = WignerDCalculator(lmax=3, dtype=torch.float32).cuda()
+        calc.eval()
+
+        saved_device = torch.get_default_device()
+        torch.set_default_device(None)
+        try:
+            with torch.device("cuda"):
+                d_eager, dt_eager = calc.forward(quat)
+                d_compiled, dt_compiled = torch.compile(calc.forward)(quat)
+            # The assembly op itself is bitwise (see the reference test);
+            # the compiled ladder reassociates the degree-block math, so the
+            # cross-route comparison holds at rounding level.
+            torch.testing.assert_close(d_compiled, d_eager, atol=1e-5, rtol=1e-5)
+            torch.testing.assert_close(dt_compiled, dt_eager, atol=1e-5, rtol=1e-5)
+        finally:
+            torch.set_default_device(saved_device)
+
+        # The inverse-layout policy is one shared exit: both routes honor
+        # the compatibility flag instead of the traced route bypassing it.
+        def _run_routes():
+            saved = torch.get_default_device()
+            torch.set_default_device(None)
+            try:
+                with torch.device("cuda"):
+                    _, dt_eager_i = calc.forward(quat)
+                    _, dt_compiled_i = torch.compile(calc.forward)(quat)
+                return dt_eager_i, dt_compiled_i
+            finally:
+                torch.set_default_device(saved)
+
+        calc._materialize_inverse_rotation = False
+        dt_view, dt_view_c = _run_routes()
+        self.assertFalse(dt_view.is_contiguous())
+        self.assertFalse(dt_view_c.is_contiguous())
+        calc._materialize_inverse_rotation = True
+        dt_dense, dt_dense_c = _run_routes()
+        self.assertTrue(dt_dense.is_contiguous())
+        self.assertTrue(dt_dense_c.is_contiguous())
+        torch.testing.assert_close(dt_dense, dt_view, atol=0.0, rtol=0.0)
+
+    @_GPU_KERNELS
+    def test_lmax_zero_and_empty_edges(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        quat = torch.randn(8, 4, device="cuda")
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+        calc0 = WignerDCalculator(lmax=0, dtype=torch.float32).cuda()
+        calc0.eval()
+        d0, dt0 = calc0.forward(quat)
+        expected = torch.zeros(8, 1, 1, device="cuda")
+        expected[:, 0, 0] = 1.0
+        self.assertTrue(torch.equal(d0, expected))
+        self.assertTrue(torch.equal(dt0, expected))
+
+        calc3 = WignerDCalculator(lmax=3, dtype=torch.float32).cuda()
+        calc3.eval()
+        d_empty, dt_empty = calc3.forward(torch.empty(0, 4, device="cuda"))
+        self.assertEqual(tuple(d_empty.shape), (0, 16, 16))
+        self.assertEqual(tuple(dt_empty.shape), (0, 16, 16))
+
+
+if __name__ == "__main__":
+    unittest.main()

From c5feab50604f2e6cf1a8fd639ca4db57fb2c05b8 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 22 Oct 2026 19:40:33 +0800
Subject: [PATCH 5/7] fix(pt): keep the Wigner-D assembly importable on torch <
 2.4

The operator registration (torch.library.custom_op / register_fake /
register_autograd) requires torch >= 2.4 while the declared minimum is
2.1, and this module is imported unconditionally through the descriptor
chain: on a 2.1-2.3 install the import itself raised AttributeError
before the eager fallback could ever run. The registration is now gated
on the API's presence; without it both the eager and the traced route
keep the autograd-through-scatter implementation, and a subprocess test
reloads the module with the API hidden to pin that behavior.

The make_fx fake-mode half of the tracing probe is now covered (a
symbolic trace selects the opaque operator, asserted on the graph code),
and the two pure-PyTorch operator tests (values/grads/gradcheck and the
opcheck schema/fake/autograd contract) no longer require CUDA, so they
run in the CPU CI.
---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py |  42 ++---
 .../pt/model/test_descriptor_sezm_triton.py   | 144 ++++++++++++++++--
 2 files changed, 156 insertions(+), 30 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index f9be1ab451..2bb27a5d62 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -78,20 +78,25 @@ def _wigner_d_assemble_impl(
     return d_full
 
 
-_wigner_d_assemble_op = torch.library.custom_op(
-    "sezm_nn::wigner_d_assemble", mutates_args=()
-)(_wigner_d_assemble_impl)
-
-
-@_wigner_d_assemble_op.register_fake
-def _(
-    blocks: list[torch.Tensor],
-    dim_full: int,
-    poly_lmin: int,
-    n_regular: int,
-) -> torch.Tensor:
-    ref = blocks[0]
-    return ref.new_zeros((ref.shape[0], dim_full, dim_full))
+_wigner_d_assemble_op: Any = None
+if hasattr(torch.library, "custom_op"):
+    # torch.library.custom_op / register_fake / register_autograd arrived in
+    # torch 2.4 while the declared minimum is 2.1: on a 2.1-2.3 install the
+    # operator stays unregistered (None) and both the eager and the traced
+    # route below keep the autograd-through-scatter implementation.
+    _wigner_d_assemble_op = torch.library.custom_op(
+        "sezm_nn::wigner_d_assemble", mutates_args=()
+    )(_wigner_d_assemble_impl)
+
+    @_wigner_d_assemble_op.register_fake
+    def _(
+        blocks: list[torch.Tensor],
+        dim_full: int,
+        poly_lmin: int,
+        n_regular: int,
+    ) -> torch.Tensor:
+        ref = blocks[0]
+        return ref.new_zeros((ref.shape[0], dim_full, dim_full))
 
 
 def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None:
@@ -121,9 +126,10 @@ def _wigner_d_assemble_backward(
     return grads, None, None, None
 
 
-_wigner_d_assemble_op.register_autograd(
-    _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup
-)
+if _wigner_d_assemble_op is not None:
+    _wigner_d_assemble_op.register_autograd(
+        _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup
+    )
 
 
 class CaseCoefficients(nn.Module):
@@ -736,7 +742,7 @@ def forward(
                 device=edge_quaternion.device,
             )
             D_full[:, 0, 0] = 1.0
-        elif tracing:
+        elif tracing and _wigner_d_assemble_op is not None:
             D_full = _wigner_d_assemble_op(
                 all_blocks,
                 self.dim_full,
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 9094065d64..ab60dcbd17 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2060,10 +2060,10 @@ class TestSeZMWignerDAssembleOp(unittest.TestCase):
 
     WIDTHS = (3, 5, 7)  # lmax = 3 ladder
 
-    def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64):
-        generator = torch.Generator(device="cuda").manual_seed(seed)
+    def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64, device="cuda"):
+        generator = torch.Generator(device=device).manual_seed(seed)
         return [
-            torch.randn(n_edge, w, w, device="cuda", dtype=dtype, generator=generator)
+            torch.randn(n_edge, w, w, device=device, dtype=dtype, generator=generator)
             for w in widths
         ]
 
@@ -2082,9 +2082,9 @@ def _reference_assemble(self, blocks, poly, dim_full, poly_lmin):
         return dense
 
     def _loss(self, dense, seed):
-        generator = torch.Generator(device="cuda").manual_seed(seed)
+        generator = torch.Generator(device=dense.device).manual_seed(seed)
         w1 = torch.randn(
-            dense.shape, device="cuda", dtype=dense.dtype, generator=generator
+            dense.shape, device=dense.device, dtype=dense.dtype, generator=generator
         )
         w2 = torch.randn(
             dense.shape, device="cuda", dtype=dense.dtype, generator=generator
@@ -2092,22 +2092,27 @@ def _loss(self, dense, seed):
         # cotangent contributions through both D and its transpose
         return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum()
 
-    @_GPU_KERNELS
+    @unittest.skipUnless(
+        hasattr(torch.library, "custom_op"),
+        "the operator registration needs torch >= 2.4",
+    )
     def test_op_matches_reference_with_grads(self):
+        """Pure-PyTorch operator math; runs wherever the CPU CI runs."""
         from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
             _wigner_d_assemble_impl,
             _wigner_d_assemble_op,
         )
 
+        device = "cuda" if _CUDA else "cpu"
         n_edge = 64
         for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)):
             with self.subTest(widths=widths, tail=poly_lmin is not None):
-                blocks = self._make_blocks(widths, n_edge, seed=5)
+                blocks = self._make_blocks(widths, n_edge, seed=5, device=device)
                 poly = None
                 n_regular = len(blocks)
                 if poly_lmin is not None:
                     poly_w = 4
-                    poly = self._make_blocks([poly_w], n_edge, seed=7)[0]
+                    poly = self._make_blocks([poly_w], n_edge, seed=7, device=device)[0]
                 dim_full = sum(widths) + 1 if poly is None else poly_lmin**2 + poly_w
                 args = (
                     blocks + ([poly] if poly is not None else []),
@@ -2159,25 +2164,30 @@ def test_op_matches_reference_with_grads(self):
                     )
                 )
 
-    @_GPU_KERNELS
+    @unittest.skipUnless(
+        hasattr(torch.library, "custom_op"),
+        "the operator registration needs torch >= 2.4",
+    )
     def test_opcheck_schema_fake_autograd(self):
+        """Schema/fake/autograd contract; pure PyTorch, CPU CI included."""
         from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
             _wigner_d_assemble_op,
         )
 
+        device = "cuda" if _CUDA else "cpu"
         for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)):
             with self.subTest(widths=widths, tail=poly_lmin is not None):
-                blocks = self._make_blocks(widths, 8, seed=9)
+                blocks = self._make_blocks(widths, 8, seed=9, device=device)
                 args_tail = []
                 n_regular = len(blocks)
                 if poly_lmin is not None:
-                    args_tail = self._make_blocks([4], 8, seed=13)
+                    args_tail = self._make_blocks([4], 8, seed=13, device=device)
                     blocks = blocks + args_tail
                 dim_full = sum(widths) + 1 if not args_tail else poly_lmin**2 + 4
                 saved = torch.get_default_device()
                 torch.set_default_device(None)
                 try:
-                    with torch.device("cuda"):
+                    with torch.device(device):
                         torch.library.opcheck(
                             _wigner_d_assemble_op,
                             (blocks, dim_full, poly_lmin or 11, n_regular),
@@ -2185,6 +2195,116 @@ def test_opcheck_schema_fake_autograd(self):
                 finally:
                     torch.set_default_device(saved)
 
+    @unittest.skipUnless(
+        hasattr(torch.library, "custom_op"),
+        "the operator registration needs torch >= 2.4",
+    )
+    def test_make_fx_trace_selects_the_assemble_op(self):
+        """The fake-mode probe half: make_fx tracing picks the opaque op.
+
+        ``torch.compile`` (covered by the CUDA test below) satisfies the
+        probe through ``is_compiling``; make_fx keeps that flag False and
+        relies on the fake-mode disjunct, which this check exercises.  The
+        graph-node assertion locks in the mechanism: with the probe
+        silently false, the traced route would fall back to the raw
+        scatter and every value comparison would still pass.
+        """
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        def _run(quat):
+            return calc.forward(quat)
+
+        calc = WignerDCalculator(lmax=3, dtype=torch.float32).cpu()
+        calc.eval()
+        quat = torch.randn(8, 4, dtype=torch.float32, device="cpu")
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+        graph = make_fx(_run, tracing_mode="symbolic", _allow_non_fake_inputs=True)(
+            quat
+        )
+        self.assertIn("sezm_nn.wigner_d_assemble", graph.code)
+
+    def test_import_and_tracing_fall_back_without_custom_op(self):
+        """A torch 2.1-2.3 install must import the module and stay eager.
+
+        Emulated in a subprocess: the module reloads with
+        ``torch.library.custom_op`` hidden for exactly the duration of its
+        top-level re-execution (torch's own lazily importing users of the
+        API stay intact), after which the operator is unregistered, the
+        eager forward/backward works, and make_fx tracing keeps the eager
+        implementation instead of the operator node.
+        """
+        import pathlib
+        import subprocess
+        import sys
+        import textwrap
+
+        import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod
+
+        tree = pathlib.Path(wignerd_mod.__file__).resolve().parents[5]
+        child = textwrap.dedent(
+            """
+            import pathlib
+            import sys
+
+            tree = pathlib.Path(sys.argv[1])
+            sys.meta_path = [
+                f
+                for f in sys.meta_path
+                if "editable"
+                not in (type(f).__module__ + "." + type(f).__name__).lower()
+                and "editable" not in getattr(f, "__name__", "").lower()
+            ]
+            sys.path.insert(0, str(tree))
+            import importlib
+
+            import torch
+            import torch.distributed._functional_collectives  # noqa: F401
+
+            import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd
+
+            assert wd._wigner_d_assemble_op is not None, "control: op registers"
+            saved_custom_op = torch.library.custom_op
+            del torch.library.custom_op
+            wd = importlib.reload(wd)
+            torch.library.custom_op = saved_custom_op
+            assert wd._wigner_d_assemble_op is None, "fallback: op stays out"
+
+            calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu()
+            calc.eval()
+            quat = torch.randn(4, 4, dtype=torch.float64, device="cpu")
+            quat = (quat / quat.norm(dim=-1, keepdim=True)).requires_grad_(True)
+            dense, _dt = calc.forward(quat)
+            dense.sum().backward()
+
+            from torch.fx.experimental.proxy_tensor import make_fx
+
+            def _run(q):
+                return calc.forward(q)
+
+            graph = make_fx(
+                _run, tracing_mode="symbolic", _allow_non_fake_inputs=True
+            )(quat)
+            assert (
+                "sezm_nn.wigner_d_assemble" not in graph.code
+            ), "without the API the traced route keeps the eager path"
+            print("fallback-ok")
+            """
+        )
+        result = subprocess.run(
+            [sys.executable, "-c", child, str(tree)],
+            capture_output=True,
+            text=True,
+            timeout=300,
+        )
+        self.assertEqual(
+            result.returncode,
+            0,
+            f"fallback child failed:\n{result.stdout}\n{result.stderr}",
+        )
+        self.assertIn("fallback-ok", result.stdout)
+
     @_GPU_KERNELS
     def test_traced_matches_eager_and_layout_policy(self):
         from deepmd.pt.model.descriptor.sezm_nn.wignerd import (

From 9aeb0b910b342dba6e3cd438d1d9f1bbdc919d43 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 25 Oct 2026 13:34:56 +0800
Subject: [PATCH 6/7] fix(pt): fall back to the scatter assembly on any
 registration failure

The schema inference for list[torch.Tensor] parameters only supports
this signature on later torches still (2.4-2.6 have the custom-op API
but reject this signature outright; 2.1-2.3 lack the API), so the
presence gate alone was insufficient. Registration now runs under a
try/except and any failure keeps the autograd-through-scatter route for
both the eager and the traced call.

The subprocess fallback test covers both incompatible classes with
separate emulations -- the API hidden entirely (2.1-2.3) and the API
replaced by a stub raising the schema-inference error (2.4-2.6) -- with
the registration control asserted only when the API exists, so the test
also holds on a genuinely old install. The CPU-run _loss helper creates
its second weight on the operand's device instead of hardcoding CUDA.
---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 22 ++++--
 .../pt/model/test_descriptor_sezm_triton.py   | 76 ++++++++++++-------
 2 files changed, 64 insertions(+), 34 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index 2bb27a5d62..42a72e557a 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -79,11 +79,14 @@ def _wigner_d_assemble_impl(
 
 
 _wigner_d_assemble_op: Any = None
-if hasattr(torch.library, "custom_op"):
+try:
     # torch.library.custom_op / register_fake / register_autograd arrived in
-    # torch 2.4 while the declared minimum is 2.1: on a 2.1-2.3 install the
-    # operator stays unregistered (None) and both the eager and the traced
-    # route below keep the autograd-through-scatter implementation.
+    # torch 2.4, and the schema inference for ``list[torch.Tensor]``
+    # parameters only supports this signature on later releases still
+    # (torch 2.4-2.6 reject it outright; 2.1-2.3 lack the API).  The
+    # declared minimum is 2.1, so any registration failure keeps the
+    # autograd-through-scatter implementation for both the eager and the
+    # traced route below instead of failing the import.
     _wigner_d_assemble_op = torch.library.custom_op(
         "sezm_nn::wigner_d_assemble", mutates_args=()
     )(_wigner_d_assemble_impl)
@@ -97,6 +100,8 @@ def _(
     ) -> torch.Tensor:
         ref = blocks[0]
         return ref.new_zeros((ref.shape[0], dim_full, dim_full))
+except Exception:  # pragma: no cover - depends on the installed torch
+    _wigner_d_assemble_op = None
 
 
 def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None:
@@ -127,9 +132,12 @@ def _wigner_d_assemble_backward(
 
 
 if _wigner_d_assemble_op is not None:
-    _wigner_d_assemble_op.register_autograd(
-        _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup
-    )
+    try:
+        _wigner_d_assemble_op.register_autograd(
+            _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup
+        )
+    except Exception:  # pragma: no cover - depends on the installed torch
+        _wigner_d_assemble_op = None
 
 
 class CaseCoefficients(nn.Module):
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index ab60dcbd17..fc46922bcd 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2087,7 +2087,7 @@ def _loss(self, dense, seed):
             dense.shape, device=dense.device, dtype=dense.dtype, generator=generator
         )
         w2 = torch.randn(
-            dense.shape, device="cuda", dtype=dense.dtype, generator=generator
+            dense.shape, device=dense.device, dtype=dense.dtype, generator=generator
         )
         # cotangent contributions through both D and its transpose
         return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum()
@@ -2225,14 +2225,18 @@ def _run(quat):
         )
         self.assertIn("sezm_nn.wigner_d_assemble", graph.code)
 
-    def test_import_and_tracing_fall_back_without_custom_op(self):
-        """A torch 2.1-2.3 install must import the module and stay eager.
-
-        Emulated in a subprocess: the module reloads with
-        ``torch.library.custom_op`` hidden for exactly the duration of its
-        top-level re-execution (torch's own lazily importing users of the
-        API stay intact), after which the operator is unregistered, the
-        eager forward/backward works, and make_fx tracing keeps the eager
+    def test_import_and_tracing_fall_back_when_registration_fails(self):
+        """Unsupported torches must import the module and stay eager.
+
+        Two emulations run in subprocesses, one per incompatible class:
+        ``hide`` deletes ``torch.library.custom_op`` for the module's
+        re-execution (torch 2.1-2.3, where the API is absent), and
+        ``reject`` replaces it with a stub raising the schema-inference
+        error (torch 2.4-2.6, where the API exists but rejects a
+        ``list[torch.Tensor]`` parameter). In both, torch's own lazily
+        importing users of the API stay intact (the change is scoped to
+        the reload window), the operator stays unregistered, the eager
+        forward/backward works, and make_fx tracing keeps the eager
         implementation instead of the operator node.
         """
         import pathlib
@@ -2248,7 +2252,8 @@ def test_import_and_tracing_fall_back_without_custom_op(self):
             import pathlib
             import sys
 
-            tree = pathlib.Path(sys.argv[1])
+            mode = sys.argv[1]
+            tree = pathlib.Path(sys.argv[2])
             sys.meta_path = [
                 f
                 for f in sys.meta_path
@@ -2264,11 +2269,25 @@ def test_import_and_tracing_fall_back_without_custom_op(self):
 
             import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd
 
-            assert wd._wigner_d_assemble_op is not None, "control: op registers"
-            saved_custom_op = torch.library.custom_op
-            del torch.library.custom_op
+            api_present = hasattr(torch.library, "custom_op")
+            if api_present:
+                assert wd._wigner_d_assemble_op is not None, "control: registers"
+
+            saved_custom_op = getattr(torch.library, "custom_op", None)
+            if mode == "hide":
+                del torch.library.custom_op
+            elif mode == "reject":
+                def _rejected(*args, **kwargs):
+                    raise TypeError(
+                        "cannot infer schema for blocks: list[torch.Tensor]"
+                    )
+
+                torch.library.custom_op = _rejected
+            else:
+                raise AssertionError(f"unknown mode {mode}")
             wd = importlib.reload(wd)
-            torch.library.custom_op = saved_custom_op
+            if saved_custom_op is not None:
+                torch.library.custom_op = saved_custom_op
             assert wd._wigner_d_assemble_op is None, "fallback: op stays out"
 
             calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu()
@@ -2288,22 +2307,25 @@ def _run(q):
             )(quat)
             assert (
                 "sezm_nn.wigner_d_assemble" not in graph.code
-            ), "without the API the traced route keeps the eager path"
+            ), "without a registration the traced route keeps the eager path"
             print("fallback-ok")
             """
         )
-        result = subprocess.run(
-            [sys.executable, "-c", child, str(tree)],
-            capture_output=True,
-            text=True,
-            timeout=300,
-        )
-        self.assertEqual(
-            result.returncode,
-            0,
-            f"fallback child failed:\n{result.stdout}\n{result.stderr}",
-        )
-        self.assertIn("fallback-ok", result.stdout)
+        for mode in ("hide", "reject"):
+            with self.subTest(emulation=mode):
+                result = subprocess.run(
+                    [sys.executable, "-c", child, mode, str(tree)],
+                    capture_output=True,
+                    text=True,
+                    timeout=300,
+                )
+                self.assertEqual(
+                    result.returncode,
+                    0,
+                    f"fallback child ({mode}) failed:"
+                    f"\n{result.stdout}\n{result.stderr}",
+                )
+                self.assertIn("fallback-ok", result.stdout)
 
     @_GPU_KERNELS
     def test_traced_matches_eager_and_layout_policy(self):

From 64696312ac758c86be71a8e89767e97a9a30763c Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 30 Oct 2026 19:40:23 +0800
Subject: [PATCH 7/7] fix(sezm): probe torch.compiler.is_compiling before eager
 Wigner use

torch 2.1-2.2 lack the callable, so the eager fallback path itself raised
AttributeError before reaching the operator-availability check.  Resolve
the callable once at import and treat it as not-compiling when absent.

The operator tests now gate on registration success instead of
torch.library.custom_op presence (2.4-2.6 have the API but reject the
schema), the subprocess hide mode tolerates an already-absent attribute,
and the successful-registration control asserts the parent-observed
state instead of assuming API presence implies registration.
---
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 22 ++++--
 .../pt/model/test_descriptor_sezm_triton.py   | 67 ++++++++++++++-----
 2 files changed, 65 insertions(+), 24 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index 42a72e557a..62c7bfda0d 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -44,6 +44,15 @@
 
 _TORCH_RELEASE = Version(torch.__version__).release[:2]
 
+# ``torch.compiler.is_compiling`` is absent before torch 2.3.  The eager
+# Wigner path must stay runnable on those versions, so resolve the callable
+# once at import time instead of calling it unconditionally.
+_COMPILER_IS_COMPILING = getattr(
+    getattr(torch, "compiler", None),
+    "is_compiling",
+    None,
+)
+
 
 def _wigner_d_assemble_impl(
     blocks: list[torch.Tensor],
@@ -728,13 +737,14 @@ def forward(
 
         # ``is_compiling`` alone is not enough here: this module is traced
         # through make_fx fake tensors, where the flag stays False.  A fake
-        # tensor carries its mode, so probe for that instead.
-        tracing = (
-            torch.compiler.is_compiling()
-            or getattr(blocks[0], "fake_mode", None) is not None
-            if blocks
-            else False
+        # tensor carries its mode, so probe for that instead.  On torch older
+        # than 2.3 the callable does not exist and the fake-tensor probe
+        # carries the decision alone.
+        compiling = (
+            _COMPILER_IS_COMPILING() if _COMPILER_IS_COMPILING is not None else False
         )
+        fake_mode = bool(blocks) and getattr(blocks[0], "fake_mode", None) is not None
+        tracing = compiling or fake_mode
         # One assembly implementation serves both routes: the eager call
         # goes straight to the differentiable scatter (autograd flows through
         # the slice assignments), the traced call hides it behind the opaque
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index fc46922bcd..1c6ee2ca54 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -41,6 +41,9 @@
 from deepmd.pt.model.descriptor.sezm_nn.so2 import (
     DynamicRadialDegreeMixer,
 )
+from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+    _wigner_d_assemble_op as _WIGNER_ASSEMBLE_OP,
+)
 from deepmd.pt_expt.kernels.triton.sezm import (
     TRITON_AVAILABLE,
 )
@@ -60,6 +63,12 @@
 
 _CUDA = torch.cuda.is_available()
 
+# The assembly operator registers only where the running torch accepts its
+# schema (torch >= 2.7); torch 2.4-2.6 have ``torch.library.custom_op`` but
+# reject a ``list[torch.Tensor]`` parameter, so tests that call the operator
+# directly gate on registration success, not on API presence.
+_WIGNER_OP_REGISTERED = _WIGNER_ASSEMBLE_OP is not None
+
 # All SeZM Triton kernels ship together (every per-module availability flag
 # reduces to "the triton package imports"), so one gate covers every
 # GPU-kernel test class.  Dispatch- and parsing-only classes stay ungated.
@@ -2093,8 +2102,8 @@ def _loss(self, dense, seed):
         return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum()
 
     @unittest.skipUnless(
-        hasattr(torch.library, "custom_op"),
-        "the operator registration needs torch >= 2.4",
+        _WIGNER_OP_REGISTERED,
+        "the operator registration needs a torch that accepts its schema",
     )
     def test_op_matches_reference_with_grads(self):
         """Pure-PyTorch operator math; runs wherever the CPU CI runs."""
@@ -2165,8 +2174,8 @@ def test_op_matches_reference_with_grads(self):
                 )
 
     @unittest.skipUnless(
-        hasattr(torch.library, "custom_op"),
-        "the operator registration needs torch >= 2.4",
+        _WIGNER_OP_REGISTERED,
+        "the operator registration needs a torch that accepts its schema",
     )
     def test_opcheck_schema_fake_autograd(self):
         """Schema/fake/autograd contract; pure PyTorch, CPU CI included."""
@@ -2196,8 +2205,8 @@ def test_opcheck_schema_fake_autograd(self):
                     torch.set_default_device(saved)
 
     @unittest.skipUnless(
-        hasattr(torch.library, "custom_op"),
-        "the operator registration needs torch >= 2.4",
+        _WIGNER_OP_REGISTERED,
+        "the operator registration needs a torch that accepts its schema",
     )
     def test_make_fx_trace_selects_the_assemble_op(self):
         """The fake-mode probe half: make_fx tracing picks the opaque op.
@@ -2229,15 +2238,23 @@ def test_import_and_tracing_fall_back_when_registration_fails(self):
         """Unsupported torches must import the module and stay eager.
 
         Two emulations run in subprocesses, one per incompatible class:
-        ``hide`` deletes ``torch.library.custom_op`` for the module's
-        re-execution (torch 2.1-2.3, where the API is absent), and
-        ``reject`` replaces it with a stub raising the schema-inference
-        error (torch 2.4-2.6, where the API exists but rejects a
-        ``list[torch.Tensor]`` parameter). In both, torch's own lazily
-        importing users of the API stay intact (the change is scoped to
-        the reload window), the operator stays unregistered, the eager
+        ``hide`` removes ``torch.library.custom_op`` for the module's
+        re-execution (torch 2.1-2.3, where the API is absent -- on a torch
+        that really lacks it the attribute is already gone and nothing is
+        deleted), and ``reject`` replaces it with a stub raising the
+        schema-inference error (torch 2.4-2.6, where the API exists but
+        rejects a ``list[torch.Tensor]`` parameter). In both, torch's own
+        lazily importing users of the API stay intact (the change is scoped
+        to the reload window), the operator stays unregistered, the eager
         forward/backward works, and make_fx tracing keeps the eager
         implementation instead of the operator node.
+
+        The child first imports the module untouched and asserts the
+        registration state its parent observed (passed as an argument), so
+        the successful-registration control is measured against reality
+        instead of API presence: on a torch whose registration legitimately
+        fails at import, the control asserts the import-time fallback, and
+        the operator-specific tests above are skipped by their decorator.
         """
         import pathlib
         import subprocess
@@ -2254,6 +2271,7 @@ def test_import_and_tracing_fall_back_when_registration_fails(self):
 
             mode = sys.argv[1]
             tree = pathlib.Path(sys.argv[2])
+            expected_registered = sys.argv[3] == "1"
             sys.meta_path = [
                 f
                 for f in sys.meta_path
@@ -2269,13 +2287,15 @@ def test_import_and_tracing_fall_back_when_registration_fails(self):
 
             import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd
 
-            api_present = hasattr(torch.library, "custom_op")
-            if api_present:
+            if expected_registered:
                 assert wd._wigner_d_assemble_op is not None, "control: registers"
+            else:
+                assert wd._wigner_d_assemble_op is None, "control: import fallback"
 
             saved_custom_op = getattr(torch.library, "custom_op", None)
             if mode == "hide":
-                del torch.library.custom_op
+                if saved_custom_op is not None:
+                    del torch.library.custom_op
             elif mode == "reject":
                 def _rejected(*args, **kwargs):
                     raise TypeError(
@@ -2288,6 +2308,9 @@ def _rejected(*args, **kwargs):
             wd = importlib.reload(wd)
             if saved_custom_op is not None:
                 torch.library.custom_op = saved_custom_op
+            elif mode == "reject":
+                # The stub was injected over a missing attribute; remove it.
+                del torch.library.custom_op
             assert wd._wigner_d_assemble_op is None, "fallback: op stays out"
 
             calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu()
@@ -2311,10 +2334,18 @@ def _run(q):
             print("fallback-ok")
             """
         )
+        registered = wignerd_mod._wigner_d_assemble_op is not None
         for mode in ("hide", "reject"):
-            with self.subTest(emulation=mode):
+            with self.subTest(emulation=mode, registered=registered):
                 result = subprocess.run(
-                    [sys.executable, "-c", child, mode, str(tree)],
+                    [
+                        sys.executable,
+                        "-c",
+                        child,
+                        mode,
+                        str(tree),
+                        "1" if registered else "0",
+                    ],
                     capture_output=True,
                     text=True,
                     timeout=300,