docs(dpmodel,pt_expt): numpydoc sections for all remaining PR-added methods

Audit of every method added by this branch against the package docstring
convention (numpydoc with underlined Parameters/Returns/Raises sections)
found 13 prose-only or section-incomplete docstrings; this completes them:

- repformers.py: _cal_hg_graph, _cal_grrg_graph, symmetrization_op_graph,
  _exchange_ghosts_graph, Atten2MultiHeadApply.call_graph,
  Atten2EquiVarApply.call_graph (full Parameters/Returns);
  _update_g1_conv_graph, _update_g2_g1g1_graph (Returns added).
- dpa2.py: uses_graph_lower (Returns), _block_graph (comment block converted
  to a full docstring; also updates the stale slice rationale -- the
  phantom-count compensation made the softmax denominator width-independent,
  the slice is kept for the minimal static pair enumeration).
- pt_expt: repformers._exchange_ghosts_graph (Parameters/Returns/Raises),
  training._model_trace_device (Parameters/Returns),
  serialization._build_graph_dynamic_shapes_with_comm (Returns).

The inner export closure fn() in forward_lower_graph_exportable_with_comm
intentionally stays docstring-less, matching the file's two pre-existing fn
closures (implementation detail of the documented enclosing method).

Author: Han Wang

deepmd/dpmodel/descriptor/dpa2.py

@@ -719,6 +719,12 @@ def uses_graph_lower(self) -> bool:
         (needs the angle machinery -- PR-G-dpa3), compressed descriptors
         (geo/tebd tabulation is dense-only), and the explicit disable flag
         (used by e.g. the spin model wrapper).
+
+        Returns
+        -------
+        uses_graph_lower : bool
+            Whether the graph-native lower (:meth:`call_graph`) is supported
+            for the current descriptor configuration.
         """
         if self._graph_lower_disabled:
             return False
@@ -1124,27 +1130,45 @@ def call_graph(
         e_ax = graph.edge_mask.shape[0]
 
         def _block_graph(rc: float, ns: int) -> tuple[Any, int | None]:
-            # graph analogue of build_multiple_neighbor_list (nlist.py:408):
-            # dist mask always; slot TRUNCATION ONLY in the shape-static
-            # dense-adapter layout, replicating the dense
-            # `nlist[:, :, :ns]` slicing.
-            #
-            # This must be a genuine array-width SLICE, not just an
-            # edge_mask AND: the segment_sum-based channels only see zero
-            # contributions from masked-out padding regardless of the
-            # array's width, but the smooth-attention softmax
-            # (RepformerLayer.call_graph) keeps every padding PAIR in the
-            # denominator at exp(-attnw_shift) (dpa1 precedent) -- so the
-            # padding-pair COUNT, governed by the static width handed to
-            # `center_edge_pairs`/`static_nnei` and not merely by the mask
-            # contents, must match the dense sub-nlist width `ns` bit-for-
-            # bit, or the attention normalization silently drifts by
-            # O(1e-4) (extra always-masked pairs still contribute
-            # exp(-shift) to the denominator). Carry-all graphs
-            # (static_nnei is None) have no static width at all: sel is
-            # normalization-only there (spec decision #9), and the compact
-            # attention pairing groups per-center dynamically, so no
-            # truncation is needed or possible.
+            """Per-block view of the model graph (local closure).
+
+            Graph analogue of ``build_multiple_neighbor_list``
+            (nlist.py:408): dist mask always; slot TRUNCATION only in the
+            shape-static dense-adapter layout, replicating the dense
+            ``nlist[:, :, :ns]`` slicing.
+
+            Parameters
+            ----------
+            rc
+                The block cutoff radius.
+            ns
+                The block sel (dense sub-nlist width).
+
+            Returns
+            -------
+            block_graph : NeighborGraph
+                The block-restricted graph view.
+            block_static_nnei : int | None
+                The block's static width (``ns``) in the shape-static
+                layout, ``None`` for carry-all graphs.
+
+            Notes
+            -----
+            The genuine array-width SLICE (rather than an edge_mask AND)
+            keeps the shape-static pair enumeration
+            (``center_edge_pairs``/``static_nnei``) at exactly the dense
+            sub-nlist width ``ns``: the pair count stays minimal, and the
+            attention softmax sees the same present-slot layout as the
+            dense body. (Before the fixed-phantom-count compensation in
+            ``segment_softmax`` this width-match was also required for
+            numeric parity -- extra always-masked pairs each contributed
+            ``exp(-attnw_shift)`` to the denominator; the compensation has
+            since made the denominator width-independent.) Carry-all graphs
+            (``static_nnei is None``) have no static width at all: sel is
+            normalization-only there (spec decision #9), and the compact
+            attention pairing groups per-center dynamically, so no
+            truncation is needed or possible.
+            """
             if static_nnei is None or ns >= static_nnei:
                 m = graph.edge_mask & (dist <= rc)
                 return dataclasses.replace(graph, edge_mask=m), static_nnei

deepmd/dpmodel/descriptor/repformers.py

@@ -540,6 +540,27 @@ def _exchange_ghosts_graph(
         pt_expt subclass overrides this to overwrite halo rows via
         ``deepmd_export::border_op`` when ``comm_dict`` is provided (C++
         multi-rank extended-region graphs).
+
+        Parameters
+        ----------
+        g1
+            Flat node-wise atomic invariant rep, with shape [n_total, ng1].
+        comm_dict
+            MPI communication metadata; must be ``None`` on this (dpmodel)
+            path.
+        n_total
+            Total number of nodes (unused here; the pt_expt override needs
+            it).
+
+        Returns
+        -------
+        g1 : Array
+            The (unchanged) node channel, with shape [n_total, ng1].
+
+        Raises
+        ------
+        NotImplementedError
+            If ``comm_dict`` is not ``None``.
         """
         del n_total
         if comm_dict is not None:
@@ -1124,8 +1145,34 @@ def _cal_hg_graph(
 ) -> Array:
     """Graph twin of :func:`_cal_hg`: hg[n, a, b] = sum_{e: dst(e)=n} h_e[a] g_e[b].
 
-    ``nnei`` is the block sel (the smooth-branch normalization constant, spec
-    decision #9: sel = normalization only).
+    Parameters
+    ----------
+    g
+        Flat edge-wise invariant rep, with shape [n_edge, ng].
+    h
+        Flat edge-wise equivariant rep, with shape [n_edge, 3].
+    edge_mask
+        Edge mask, where zero means no edge, with shape [n_edge].
+    sw
+        The switch function per edge, with shape [n_edge].
+    dst
+        Destination (center) node index of each edge, with shape [n_edge].
+    n_total
+        Total number of nodes.
+    nnei
+        The block sel (the smooth-branch normalization constant, spec
+        decision #9: sel = normalization only).
+    smooth
+        Whether to use the smooth (sw-weighted, sel-normalized) branch.
+    epsilon
+        Degree-normalization protection for the non-smooth branch.
+    use_sqrt_nnei
+        Whether to normalize by sqrt(nnei) instead of nnei.
+
+    Returns
+    -------
+    hg : Array
+        Transposed rotation matrix per node, with shape [n_total, 3, ng].
     """
     from deepmd.dpmodel.utils.neighbor_graph import (
         segment_sum,
@@ -1148,7 +1195,20 @@ def _cal_hg_graph(
 
 
 def _cal_grrg_graph(hg: Array, axis_neuron: int) -> Array:
-    """Graph twin of :func:`_cal_grrg` (node-local, no neighbor axis)."""
+    """Graph twin of :func:`_cal_grrg` (node-local, no neighbor axis).
+
+    Parameters
+    ----------
+    hg
+        Transposed rotation matrix per node, with shape [n_total, 3, ng].
+    axis_neuron
+        Size of the submatrix of hg (embedding matrix).
+
+    Returns
+    -------
+    grrg : Array
+        Atomic invariant rep, with shape [n_total, axis_neuron * ng].
+    """
     xp = array_api_compat.array_namespace(hg)
     n, _, ng = hg.shape
     hgm = hg[..., :axis_neuron]  # (N, 3, axis)
@@ -1169,7 +1229,38 @@ def symmetrization_op_graph(
     epsilon: float = 1e-4,
     use_sqrt_nnei: bool = True,
 ) -> Array:
-    """Graph twin of :func:`symmetrization_op`."""
+    """Graph twin of :func:`symmetrization_op`.
+
+    Parameters
+    ----------
+    g
+        Flat edge-wise invariant rep, with shape [n_edge, ng].
+    h
+        Flat edge-wise equivariant rep, with shape [n_edge, 3].
+    edge_mask
+        Edge mask, where zero means no edge, with shape [n_edge].
+    sw
+        The switch function per edge, with shape [n_edge].
+    dst
+        Destination (center) node index of each edge, with shape [n_edge].
+    n_total
+        Total number of nodes.
+    nnei
+        The block sel (the smooth-branch normalization constant).
+    axis_neuron
+        Size of the submatrix of hg (embedding matrix).
+    smooth
+        Whether to use the smooth (sw-weighted, sel-normalized) branch.
+    epsilon
+        Degree-normalization protection for the non-smooth branch.
+    use_sqrt_nnei
+        Whether to normalize by sqrt(nnei) instead of nnei.
+
+    Returns
+    -------
+    grrg : Array
+        Atomic invariant rep, with shape [n_total, axis_neuron * ng].
+    """
     hg = _cal_hg_graph(
         g,
         h,
@@ -1480,7 +1571,26 @@ def call(
     def call_graph(
         self, AA: Array, g2: Array, q_e: Array, k_e: Array, e_tot: int
     ) -> Array:
-        """Graph twin of :meth:`call`: out[q] = sum_k AA[q,k] g2v[k] per head."""
+        """Graph twin of :meth:`call`: out[q] = sum_k AA[q,k] g2v[k] per head.
+
+        Parameters
+        ----------
+        AA
+            Attention map on the flat pair axis, with shape [n_pair, nh].
+        g2
+            Flat edge-wise pair invariant rep, with shape [n_edge, ng2].
+        q_e
+            Query edge index of each pair, with shape [n_pair].
+        k_e
+            Key edge index of each pair, with shape [n_pair].
+        e_tot
+            Total number of edges.
+
+        Returns
+        -------
+        g2_new : Array
+            Attention-updated edge channel, with shape [n_edge, ng2].
+        """
         from deepmd.dpmodel.utils.neighbor_graph import (
             segment_sum,
         )
@@ -1578,7 +1688,27 @@ def call(
     def call_graph(
         self, AA: Array, h2: Array, q_e: Array, k_e: Array, e_tot: int
     ) -> Array:
-        """Graph twin of :meth:`call` (heads applied to the equivariant channel)."""
+        """Graph twin of :meth:`call` (heads applied to the equivariant channel).
+
+        Parameters
+        ----------
+        AA
+            Attention map on the flat pair axis, with shape [n_pair, nh].
+        h2
+            Flat edge-wise pair equivariant rep, with shape [n_edge, 3].
+        q_e
+            Query edge index of each pair, with shape [n_pair].
+        k_e
+            Key edge index of each pair, with shape [n_pair].
+        e_tot
+            Total number of edges.
+
+        Returns
+        -------
+        h2_new : Array
+            Attention-updated equivariant edge channel, with shape
+            [n_edge, 3].
+        """
         from deepmd.dpmodel.utils.neighbor_graph import (
             segment_sum,
         )
@@ -2260,6 +2390,11 @@ def _update_g1_conv_graph(
             The block sel: the smooth-branch normalization constant AND the
             fixed dense softmax width for the attention phantom-count
             compensation (see :meth:`Atten2Map.call_graph`).
+
+        Returns
+        -------
+        g1_conv : Array
+            Convolution-updated node channel, with shape [n_total, ng1].
         """
         from deepmd.dpmodel.utils.neighbor_graph import (
             segment_sum,
@@ -2339,6 +2474,11 @@ def _update_g2_g1g1_graph(
             Edge mask, where zero means no edge, with shape [n_edge].
         sw
             The switch function, with shape [n_edge].
+
+        Returns
+        -------
+        g1g1 : Array
+            Per-edge center-neighbor g1 product, with shape [n_edge, ng1].
         """
         xp = array_api_compat.array_namespace(g1, edge_mask, sw)
         # (n_edge, ng1)

deepmd/pt_expt/descriptor/repformers.py

@@ -108,6 +108,28 @@ def _exchange_ghosts_graph(
         (ghost-free Python graphs / extended single-process graphs).  Spin
         models never route the graph lower (``disable_graph_lower``), so a
         ``has_spin`` comm_dict reaching here is a programming error.
+
+        Parameters
+        ----------
+        g1
+            Flat node-wise atomic invariant rep, with shape [n_total, ng1].
+        comm_dict
+            MPI communication metadata (``send_list``, ``send_proc``,
+            ``recv_proc``, ``send_num``, ``recv_num``, ``communicator``,
+            ``nlocal``, ``nghost``); ``None`` for single-process graphs.
+        n_total
+            Total number of nodes (``nlocal + nghost``).
+
+        Returns
+        -------
+        g1 : torch.Tensor
+            The node channel with halo rows refreshed, with shape
+            [n_total, ng1].
+
+        Raises
+        ------
+        NotImplementedError
+            If ``comm_dict`` carries ``has_spin``.
         """
         if comm_dict is None:
             return super()._exchange_ghosts_graph(g1, comm_dict, n_total)

deepmd/pt_expt/train/training.py

@@ -627,6 +627,17 @@ def _model_trace_device(model: torch.nn.Module) -> torch.device:
     callers may legitimately hold the model on a different device (e.g. a test
     pinning the model to CPU while running on a CUDA host). Falls back to
     :data:`DEVICE` only if the model exposes no parameters or buffers.
+
+    Parameters
+    ----------
+    model
+        The model whose parameter/buffer device determines the trace device.
+
+    Returns
+    -------
+    device : torch.device
+        The device of the model's first parameter (or buffer), falling back
+        to the global :data:`DEVICE`.
     """
     for _t in model.parameters():
         return _t.device

deepmd/pt_expt/utils/serialization.py

@@ -538,6 +538,13 @@ def _build_graph_dynamic_shapes_with_comm(
         charge_spin, send_list, send_proc, recv_proc, send_num, recv_num,
         communicator, nlocal, nghost)`` — 16 entries matching
         ``forward_lower_graph_exportable_with_comm``.
+
+    Returns
+    -------
+    dynamic_shapes : tuple
+        Per-input dynamic-shape specs (dicts of ``torch.export.Dim`` or
+        ``None``) in the same order as ``sample_inputs``, for
+        ``torch.export.export(..., dynamic_shapes=...)``.
     """
     fparam = sample_inputs[5]
     aparam = sample_inputs[6]