feat(pt_expt): implement se_t and se_t_tebd descriptors. #5208
feat(pt_expt): implement se_t and se_t_tebd descriptors. #5208wanghan-iapcm wants to merge 37 commits intodeepmodeling:masterfrom
Conversation
Summary of ChangesHello @wanghan-iapcm, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly advances the PyTorch experimental backend by integrating Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces se_t and se_t_tebd descriptors for the pt_expt backend and includes a significant, well-executed refactoring. The key improvements are the adoption of array_api_compat for better backend-agnosticism, the refactoring of EmbeddingNet into a concrete class, and the introduction of a registry-based system in pt_expt/common.py for converting dpmodel objects to PyTorch modules. This new system greatly simplifies the __setattr__ logic and enhances maintainability. The new features are well-tested. Overall, this is a high-quality contribution. I have one minor suggestion for code cleanup.
| dtype=des["dstd"].dtype, | ||
| device=array_api_compat.device(des["dstd"]), | ||
| ) | ||
| xp = array_api_compat.array_namespace(des["davg"]) |
There was a problem hiding this comment.
This line re-defines xp. It was already defined on line 177 if the else branch was taken. To improve clarity and avoid redundancy, consider defining xp once before the if-else block (e.g., at line 171). This would make xp available in both branches without re-definition.
📝 WalkthroughWalkthroughThis PR migrates descriptor implementations to array_api_compat for cross-backend device support, introduces a concrete EmbeddingNet class replacing factory pattern generation, adds PyTorch experimental wrapper classes for descriptors, casts indices to int64 for PyTorch compatibility, and extends test coverage for new implementations. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@deepmd/dpmodel/utils/network.py`:
- Around line 886-893: The code assumes obj.layers has at least one element
(uses type(obj.layers[0])) which throws IndexError when neuron==[]; fix by
checking for empty layers first: if either the existing obj.layers or the
incoming layers list is empty, set obj.layers to an empty instance of the same
container type (use container_type = type(obj.layers); obj.layers =
container_type() or container_type([])) and return obj; otherwise proceed with
the existing logic that computes layer_type = type(obj.layers[0]) and
reconstructs obj.layers with [layer_type.deserialize(...)]. Ensure you reference
cls, obj.layers, layer_type, and the deserialize call when making the change.
In `@source/tests/pt_expt/descriptor/test_se_t_tebd.py`:
- Around line 62-71: The test is confusing because pt_expt's forward accepts a
type_embedding but the pt_expt implementation deletes it (del type_embedding in
se_t_tebd.forward) and dpmodel.call() generates its own via its internal
TypeEmbedNet; update the test
(source/tests/pt_expt/descriptor/test_se_t_tebd.py) by adding a short clarifying
comment immediately above the creation/use of type_embedding (near the dd0(...)
call) stating that the passed type_embedding is intentionally ignored by the
pt_expt wrapper and that all models use the internal TypeEmbedNet with identical
serialized state, so the external tensor is provided only to keep API parity for
the test.
🧹 Nitpick comments (6)
deepmd/pt_expt/utils/type_embed.py (1)
15-15: Unusednoqadirective.Ruff reports that
F401is not enabled, making the# noqa: F401directive unnecessary. The side-effect import itself is correct and well-documented by the comment above.Proposed fix
-from deepmd.pt_expt.utils import network # noqa: F401 +from deepmd.pt_expt.utils import network # ensure EmbeddingNet is registereddeepmd/pt_expt/descriptor/__init__.py (1)
2-3: Unused# noqa: F401directive.Ruff (RUF100) reports that the
F401rule is not enabled, making thenoqadirective a no-op. IfF401is indeed not active in this project's Ruff config, you can drop the directive.deepmd/dpmodel/utils/network.py (1)
813-816: Mutable default argument forneuron.Ruff B006 flags
neuron: list[int] = [24, 48, 96]. While the list isn't mutated in-place here, usingNoneas the default is the idiomatic safeguard.Suggested fix
def __init__( self, in_dim: int, - neuron: list[int] = [24, 48, 96], + neuron: list[int] | None = None, activation_function: str = "tanh", resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None, bias: bool = True, trainable: bool | list[bool] = True, ) -> None: + if neuron is None: + neuron = [24, 48, 96] layers = []source/tests/consistent/descriptor/test_se_t.py (1)
99-107: Unused unpacked variables — consistent with existing pattern, safe to ignore.Ruff flags
resnet_dt,excluded_types,precision, andenv_protectionas unused. This is consistent withskip_ptandskip_dpabove which also unpack without using the values. The pattern exists so that if future skip logic needs to use params, the unpacking is already in place (asskip_tfdoes on line 117).source/tests/pt_expt/descriptor/test_se_t_tebd.py (1)
159-159: Nit: missing trailing newline at end of file.Proposed fix
torch.export.export(dd0, inputs) +source/tests/pt_expt/descriptor/test_se_t.py (1)
65-65: Prefix unused unpacked variables with_to satisfy linter.
gr1(Line 65) andgr2(Line 85) are unpacked but never used. The comment on Line 90 confirmsse_treturnsNoneforgr, so these should be prefixed with_.Proposed fix
- rd1, gr1, _, _, sw1 = dd1( + rd1, _gr1, _, _, sw1 = dd1(- rd2, gr2, _, _, sw2 = dd2.call( + rd2, _gr2, _, _, sw2 = dd2.call(Also applies to: 85-85
| obj = cls(**data) | ||
| # Reinitialize layers from serialized data, using the same layer type | ||
| # that __init__ created (respects subclass overrides via MRO). | ||
| layer_type = type(obj.layers[0]) | ||
| obj.layers = type(obj.layers)( | ||
| [layer_type.deserialize(layer) for layer in layers] | ||
| ) | ||
| return obj |
There was a problem hiding this comment.
IndexError when neuron is an empty list.
type(obj.layers[0]) on line 889 will raise IndexError if neuron=[] (producing zero layers). While this may be an uncommon case, deserialize should handle it gracefully — especially since FittingNet already accounts for len(neuron) == 0.
Proposed fix
obj = cls(**data)
# Reinitialize layers from serialized data, using the same layer type
# that __init__ created (respects subclass overrides via MRO).
- layer_type = type(obj.layers[0])
- obj.layers = type(obj.layers)(
- [layer_type.deserialize(layer) for layer in layers]
- )
+ if obj.layers:
+ layer_type = type(obj.layers[0])
+ obj.layers = type(obj.layers)(
+ [layer_type.deserialize(layer) for layer in layers]
+ )
return obj🤖 Prompt for AI Agents
In `@deepmd/dpmodel/utils/network.py` around lines 886 - 893, The code assumes
obj.layers has at least one element (uses type(obj.layers[0])) which throws
IndexError when neuron==[]; fix by checking for empty layers first: if either
the existing obj.layers or the incoming layers list is empty, set obj.layers to
an empty instance of the same container type (use container_type =
type(obj.layers); obj.layers = container_type() or container_type([])) and
return obj; otherwise proceed with the existing logic that computes layer_type =
type(obj.layers[0]) and reconstructs obj.layers with
[layer_type.deserialize(...)]. Ensure you reference cls, obj.layers, layer_type,
and the deserialize call when making the change.
| type_embedding = torch.randn( | ||
| [self.nt, dd0.tebd_dim], dtype=dtype, device=self.device | ||
| ) | ||
|
|
||
| rd0, _, _, _, _ = dd0( | ||
| torch.tensor(self.coord_ext, dtype=dtype, device=self.device), | ||
| torch.tensor(self.atype_ext, dtype=int, device=self.device), | ||
| torch.tensor(self.nlist, dtype=int, device=self.device), | ||
| type_embedding=type_embedding, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's find the test file and understand its structure
find . -name "test_se_t_tebd.py" -type f 2>/dev/null | head -5Repository: deepmodeling/deepmd-kit
Length of output: 173
🏁 Script executed:
#!/bin/bash
# Look for the pt_expt descriptor implementation
find . -name "se_t_tebd.py" -type f 2>/dev/null | grep -E "(pt_expt|dpmodel)" | head -10Repository: deepmodeling/deepmd-kit
Length of output: 148
🏁 Script executed:
#!/bin/bash
# Search for TypeEmbedNet and how type_embedding is used in forward methods
rg -n "type_embedding" --type py -B 2 -A 5 | head -100Repository: deepmodeling/deepmd-kit
Length of output: 6671
🏁 Script executed:
#!/bin/bash
# Read the test file to see the code structure mentioned in the review
cat -n ./source/tests/pt_expt/descriptor/test_se_t_tebd.py | sed -n '55,110p'Repository: deepmodeling/deepmd-kit
Length of output: 2772
🏁 Script executed:
#!/bin/bash
# Check the pt_expt descriptor implementation
wc -l ./deepmd/pt_expt/descriptor/se_t_tebd.pyRepository: deepmodeling/deepmd-kit
Length of output: 110
🏁 Script executed:
#!/bin/bash
# Check the dpmodel descriptor implementation
wc -l ./deepmd/dpmodel/descriptor/se_t_tebd.pyRepository: deepmodeling/deepmd-kit
Length of output: 111
🏁 Script executed:
#!/bin/bash
# Read the pt_expt descriptor implementation (short file)
cat -n ./deepmd/pt_expt/descriptor/se_t_tebd.pyRepository: deepmodeling/deepmd-kit
Length of output: 2117
🏁 Script executed:
#!/bin/bash
# Search for how type_embedding is handled in dpmodel descriptor
grep -n "type_embedding\|def call\|def forward" ./deepmd/dpmodel/descriptor/se_t_tebd.py | head -30Repository: deepmodeling/deepmd-kit
Length of output: 1217
🏁 Script executed:
#!/bin/bash
# Look at the context around line 47 in pt_expt file to see if there are comments explaining the deletion
cat -n ./deepmd/pt_expt/descriptor/se_t_tebd.py | sed -n '40,54p'Repository: deepmodeling/deepmd-kit
Length of output: 594
🏁 Script executed:
#!/bin/bash
# Look at the dpmodel call() method signature to confirm it doesn't accept type_embedding
sed -n '326,380p' ./deepmd/dpmodel/descriptor/se_t_tebd.py | head -60Repository: deepmodeling/deepmd-kit
Length of output: 1945
🏁 Script executed:
#!/bin/bash
# Check if there are any comments in the test file around the assertion that might explain the intent
sed -n '90,110p' ./source/tests/pt_expt/descriptor/test_se_t_tebd.pyRepository: deepmodeling/deepmd-kit
Length of output: 839
Add a clarifying comment explaining why type_embedding is ignored in the pt_expt wrapper.
The pt_expt descriptor's forward() method accepts type_embedding as a parameter (line 39) but explicitly ignores it via del type_embedding (line 47 of ./deepmd/pt_expt/descriptor/se_t_tebd.py). It then delegates to dpmodel.call() which does not accept type_embedding and generates its own internally via self.type_embedding.call() (line 366 of ./deepmd/dpmodel/descriptor/se_t_tebd.py).
The test works because all three code paths (dd0, dd1, dd2) use the internal TypeEmbedNet with identical serialized state. However, the API is confusing: the parameter is accepted but silently discarded without documentation.
Either:
- Add a comment in the test explaining that the external
type_embeddingis intentionally ignored because both pt_expt and dpmodel use internalTypeEmbedNet, or - Remove
type_embeddingfrom the pt_exptforward()signature to avoid the misleading API.
🤖 Prompt for AI Agents
In `@source/tests/pt_expt/descriptor/test_se_t_tebd.py` around lines 62 - 71, The
test is confusing because pt_expt's forward accepts a type_embedding but the
pt_expt implementation deletes it (del type_embedding in se_t_tebd.forward) and
dpmodel.call() generates its own via its internal TypeEmbedNet; update the test
(source/tests/pt_expt/descriptor/test_se_t_tebd.py) by adding a short clarifying
comment immediately above the creation/use of type_embedding (near the dd0(...)
call) stating that the passed type_embedding is intentionally ignored by the
pt_expt wrapper and that all models use the internal TypeEmbedNet with identical
serialized state, so the external tensor is provided only to keep API parity for
the test.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 341a5edf8e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| obj = cls(**data) | ||
| # Reinitialize layers from serialized data, using the same layer type | ||
| # that __init__ created (respects subclass overrides via MRO). | ||
| layer_type = type(obj.layers[0]) |
There was a problem hiding this comment.
EmbeddingNet.deserialize() now unconditionally reads obj.layers[0], which raises IndexError when deserializing a valid zero-hidden-layer embedding (neuron=[], serialized as "layers": []). This regresses previous behavior where empty-layer embeddings could deserialize, and it can break model loading/type-conversion paths that rely on round-tripping these minimal networks.
Useful? React with 👍 / 👎.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ) | ||
|
|
||
| # Import network to ensure EmbeddingNet is registered before TypeEmbedNet is used | ||
| from deepmd.pt_expt.utils import network # noqa: F401 |
Check notice
Code scanning / CodeQL
Unused import Note
| @property | ||
| def skip_pt_expt(self) -> bool: | ||
| ( | ||
| resnet_dt, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| def skip_pt_expt(self) -> bool: | ||
| ( | ||
| resnet_dt, | ||
| excluded_types, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| ( | ||
| resnet_dt, | ||
| excluded_types, | ||
| precision, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| resnet_dt, | ||
| excluded_types, | ||
| precision, | ||
| env_protection, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| tebd_dim, | ||
| tebd_input_mode, | ||
| resnet_dt, | ||
| excluded_types, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| precision, | ||
| use_econf_tebd, | ||
| use_tebd_bias, | ||
| ) = self.param |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| use_econf_tebd, | ||
| use_tebd_bias, | ||
| ) = self.param | ||
| return CommonTest.skip_pt_expt |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| use_tebd_bias, | ||
| ) = self.param | ||
| return CommonTest.skip_pt_expt | ||
|
|
Check notice
Code scanning / CodeQL
Unused local variable Note test
| ) = self.param | ||
| return CommonTest.skip_pt_expt | ||
|
|
||
| @property |
Check notice
Code scanning / CodeQL
Unused local variable Note test
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5208 +/- ##
==========================================
+ Coverage 82.00% 82.02% +0.02%
==========================================
Files 724 728 +4
Lines 73801 73920 +119
Branches 3616 3615 -1
==========================================
+ Hits 60520 60635 +115
- Misses 12120 12121 +1
- Partials 1161 1164 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
2 participants
This PR is considered after #5194 #5204 and #5205
Summary by CodeRabbit
New Features
Bug Fixes
Refactor
Tests