deepmd/pt/model/atomic_model/linear_atomic_model.py (3)

238-257: Docstring is missing the new comm_dict parameter

forward_atomic added a comm_dict argument but the docstring hasn’t been updated. Please document it to avoid confusion for users and downstream bindings.

Apply this diff inside the existing docstring’s “Parameters” section:

         aparam
             atomic parameter. (nframes, nloc, nda)
 
+        comm_dict
+            Optional dict[str, torch.Tensor]. A scratch communication dictionary
+            forwarded to sub-models’ forward_common_atomic. Implementations may
+            read/write entries to share intermediates; pass None to disable.

---

282-295: Consider namespacing comm_dict per sub-model to avoid key collisions

If multiple sub-models write the same keys, they can clobber each other. If shared keys are not intended, namespace the dict per sub-model index. If shared keys are intended, ignore this.

Example change (only if you observe collisions in practice):

         for i, model in enumerate(self.models):
             type_map_model = self.mapping_list[i].to(extended_atype.device)
             # apply bias to each individual model
             ener_list.append(
                 model.forward_common_atomic(
                     extended_coord,
                     type_map_model[extended_atype],
                     nlists_[i],
                     mapping,
                     fparam,
                     aparam,
-                    comm_dict=comm_dict,
+                    comm_dict=(
+                        None
+                        if comm_dict is None
+                        else comm_dict.setdefault(f"linear[{i}]", {})
+                    ),
                 )["energy"]
             )

To decide, please verify whether sub-models rely on shared comm_dict keys or not.

---

226-235: TorchScript typing: prefer typing.Dict for broader compatibility

The annotation uses Optional[dict[str, torch.Tensor]]. Older TorchScript versions have been finicky with PEP 585 generics. If your CI still scripts models under older PyTorch, consider switching to Optional[Dict[str, torch.Tensor]] (from typing) across the codebase for forward_* methods.

Would you like me to prepare a follow-up patch that replaces dict[...] with Dict[...] and adds the necessary imports consistently?

deepmd/pt/model/model/dp_linear_model.py (2)

90-112: Add a brief docstring entry for comm_dict to forward_lower

Helps users of the TorchScript-exported API understand the new parameter.

     @torch.jit.export
     def forward_lower(
         self,
         extended_coord,
         extended_atype,
         nlist,
         mapping: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
     ):
+        """
+        Lower-level forward.
+        Parameters
+        ----------
+        comm_dict
+            Optional dict[str, torch.Tensor] used as a communication scratchpad
+            and forwarded to forward_common_lower.
+        """
         model_ret = self.forward_common_lower(

---

90-112: TorchScript typing consistency

Same note as in the atomic model: if CI scripts modules with an older PyTorch, consider Optional[Dict[str, torch.Tensor]] for comm_dict and import Dict from typing.

deepmd/pt/model/model/polar_model.py (3)

87-97: Use named argument for mapping for consistency

Other files (e.g., dp_linear_model.py, dp_zbl_model.py) pass mapping as a named argument. Aligning improves readability and resilience to parameter order changes.

         model_ret = self.forward_common_lower(
             extended_coord,
             extended_atype,
             nlist,
-            mapping,
+            mapping=mapping,
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
             comm_dict=comm_dict,
             extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )

---

75-86: Optional: document comm_dict in forward_lower

If you prefer method-level docs, mirror a short note about comm_dict here as well.

---

75-86: TorchScript typing consistency

Consider using Optional[Dict[str, torch.Tensor]] for broader TorchScript compatibility, as noted in other files.

deepmd/pt/model/model/dos_model.py (3)

93-103: Use named argument for mapping for consistency

Match the style used elsewhere by naming mapping explicitly.

         model_ret = self.forward_common_lower(
             extended_coord,
             extended_atype,
             nlist,
-            mapping,
+            mapping=mapping,
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
             comm_dict=comm_dict,
             extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )

---

82-92: Optional: document comm_dict in forward_lower

Add a brief note/docstring for comm_dict to aid users.

---

82-92: TorchScript typing consistency

Consider Optional[Dict[str, torch.Tensor]] to maximize TorchScript compatibility, mirroring other files.

deepmd/pt/model/model/dp_zbl_model.py (2)

90-101: Optional: document comm_dict in forward_lower

Consider a short docstring snippet for the new parameter for parity with other forward APIs.

---

90-101: TorchScript typing consistency

As with other files, consider switching to Optional[Dict[str, torch.Tensor]] if your TorchScript environment benefits from typing.Dict.

deepmd/pt/model/model/dipole_model.py (3)

101-103: Propagating comm_dict looks correct; adjust type hint for TorchScript compatibility.

Good: The new comm_dict is threaded into forward_common_lower and keeps the parameter optional, preserving call-site compatibility.

Risk: @torch.jit.export functions can be picky about annotations. Using dict[str, torch.Tensor] may not be accepted by TorchScript in all environments; typing.Dict[str, torch.Tensor] is safer and consistent across PyTorch versions.

Suggested change:

• Switch to Dict[str, torch.Tensor]
• Import Dict

Apply this diff:

@@
-from typing import (
-    Optional,
-)
+from typing import (
+    Optional,
+    Dict,
+)
@@
-        comm_dict: Optional[dict[str, torch.Tensor]] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,

Verification checklist:

• Ensure the definition of forward_common_lower accepts comm_dict with a compatible type.
• Run a TorchScript export path for DipoleModel to confirm no schema/type errors occur when scripting or tracing.
• Confirm other models in this PR use the same Dict[...] convention for consistency.

Also applies to: 111-113

---

91-103: Consider adding return type annotation and brief docstring for comm_dict.

For parity with forward(), annotate forward_lower’s return type and document the expected keys in comm_dict (e.g., required/optional keys, shapes, device). This aids users and static tooling.

Example:

-    def forward_lower(
+    def forward_lower(
         self,
         extended_coord,
         extended_atype,
         nlist,
         mapping: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
-        comm_dict: Optional[dict[str, torch.Tensor]] = None,
-    ):
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Args:
+            comm_dict: Optional dictionary for inter-stage communication.
+                Expected keys (if any): e.g. "neighbor_mask", "env_info".
+                Values are torch.Tensors on the same device as inputs.
+        """

---

55-89: API symmetry: does forward() also need comm_dict?

If bug fixes rely on comm_dict only in the lower path, this is fine. If upstream callers sometimes only use forward(), consider optionally adding comm_dict there for consistency across models. Otherwise, document that comm_dict is a lower-path-only feature.

Please confirm whether the other models updated in this PR expose comm_dict only in forward_lower() or also in forward(), and that callers don’t need it at the higher level in dipole inference/training flows.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

deepmd/pt/model/atomic_model/linear_atomic_model.py (1)

282-295: Good: comm_dict is correctly propagated to sub-models

Passing comm_dict through to each sub-model’s forward_common_atomic aligns the linear composition path with the rest of the models and should resolve missing cross-module communication for issue #4906. No functional regressions spotted around weighting or aggregation.

deepmd/pt/model/model/dp_linear_model.py (2)

100-111: Good: comm_dict added and correctly threaded to forward_common_lower

Signature change is backward-compatible (new trailing optional arg) and the value is forwarded as expected.

---

90-112: All comm_dict propagation checks out—no changes needed

I’ve verified across the PyTorch backend that:

• Every forward_lower definition includes the comm_dict: Optional[...] parameter.
• Inside each forward_lower, the call to self.forward_common_lower(...) explicitly passes comm_dict=comm_dict.
• All forward_common_atomic definitions include comm_dict in their signatures.
• Within the lower‐level implementations, forward_common_atomic is always invoked with comm_dict when called from the atomic‐aware paths.

Since the propagation of comm_dict is consistent and complete, no further modifications are required here.

deepmd/pt/model/model/polar_model.py (1)

85-97: Good: comm_dict added and forwarded to forward_common_lower

Change is localized and preserves backward compatibility.

deepmd/pt/model/model/dos_model.py (1)

91-103: Good: comm_dict added and forwarded to forward_common_lower

Change looks correct and non-breaking.

deepmd/pt/model/model/dp_zbl_model.py (1)

100-112: Good: comm_dict threaded through the ZBL lower path

Signature and forwarding look correct; backward compatibility is preserved.

deepmd/pt/model/model/dipole_model.py (1)

91-129: Ensure comm_dict Is Uniformly Plumbed Through Lower‐Level Forwards

To avoid any breaking changes when propagating the new comm_dict argument, please verify across all model implementations:

• Confirm every forward_common_lower definition includes the comm_dict: Optional[dict[str, torch.Tensor]] = None parameter in its signature.
• Confirm each forward_lower method signature likewise declares comm_dict and passes it into its call to forward_common_lower.
• Search for any external or legacy call sites of forward_lower and forward_common_lower (including in tests) that might still rely on the previous signatures—especially multi‐line signatures or invocations—and update them to pass comm_dict explicitly.

Recommended checks (run from the repo root):

# Verify comm_dict in all forward_common_lower definitions
rg -U -P 'def\s+forward_common_lower.*comm_dict' -n --type=py

# Verify comm_dict in all forward_lower definitions
rg -U -P 'def\s+forward_lower.*comm_dict' -n --type=py

# Find any calls to forward_common_lower without comm_dict
rg -U -P 'forward_common_lower\s*\(' -n --type=py | grep -v 'comm_dict'

# Find any external calls to forward_lower (e.g., in tests) that lack the comm_dict keyword
rg -U -P 'forward_lower\s*\(' -n --type=py | grep -v 'comm_dict'

Please run these checks to ensure no model variant or test is accidentally omitted.