diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index b283e3a7e5..42a72e557a 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -45,6 +45,101 @@
 _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,
+) -> 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.  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]
+    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
+
+
+_wigner_d_assemble_op: Any = None
+try:
+    # torch.library.custom_op / register_fake / register_autograd arrived in
+    # 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)
+
+    @_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))
+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:
+    blocks, _, poly_lmin, n_regular = inputs
+    # 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
+
+
+def _wigner_d_assemble_backward(
+    ctx: Any, grad_d_full: torch.Tensor
+) -> tuple[list[torch.Tensor], None, None, None]:
+    # 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 size in ctx.block_widths:
+        grads.append(grad_d_full[:, off : off + size, off : off + size])
+        off += size
+    if ctx.has_poly_tail:
+        po = ctx.poly_lmin * ctx.poly_lmin
+        grads.append(grad_d_full[:, po:, po:])
+    return grads, None, None, None
+
+
+if _wigner_d_assemble_op is not None:
+    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):
     """
     Polynomial tables for one magnitude-ordered branch of the quaternion Wigner path.
@@ -546,62 +641,64 @@ 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 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
 
         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 +713,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,16 +725,54 @@ def forward(
                     lmax=self.lmax,
                     lmin=self.poly_lmin,
                 )
-                D_full[:, self.poly_offset :, self.poly_offset :] = D_poly
 
-        # Consumers address the inverse rotation through explicit strides or
-        # PyTorch strided operators, so the transpose can share D_full's storage.
+        # ``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
+        )
+        # 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 and _wigner_d_assemble_op is not None:
+            D_full = _wigner_d_assemble_op(
+                all_blocks,
+                self.dim_full,
+                self.poly_lmin,
+                len(blocks),
+            )
+        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
 
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f359f073d6..fc46922bcd 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2053,3 +2053,351 @@ 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, device="cuda"):
+        generator = torch.Generator(device=device).manual_seed(seed)
+        return [
+            torch.randn(n_edge, w, w, device=device, 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=dense.device).manual_seed(seed)
+        w1 = torch.randn(
+            dense.shape, device=dense.device, dtype=dense.dtype, generator=generator
+        )
+        w2 = torch.randn(
+            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()
+
+    @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, 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, 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 []),
+                    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,
+                    )
+                )
+
+    @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, device=device)
+                args_tail = []
+                n_regular = len(blocks)
+                if poly_lmin is not None:
+                    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(device):
+                        torch.library.opcheck(
+                            _wigner_d_assemble_op,
+                            (blocks, dim_full, poly_lmin or 11, n_regular),
+                        )
+                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_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
+        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
+
+            mode = sys.argv[1]
+            tree = pathlib.Path(sys.argv[2])
+            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
+
+            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)
+            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()
+            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 a registration the traced route keeps the eager path"
+            print("fallback-ok")
+            """
+        )
+        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 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()