fix(tf): fix compress suffix in DescrptDPA1Compat #4243
Conversation
📝 WalkthroughWalkthroughThe changes in this pull request involve modifications to the Changes
Assessment against linked issues
Possibly related PRs
Suggested labels
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![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
🧹 Outside diff range and nitpick comments (5)
source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py (2)
47-59: Consider adding error handling for session operations.While the function works correctly, it might benefit from try-catch blocks around session operations to handle potential TensorFlow runtime errors gracefully.
Example enhancement:
def build_eval_tf(sess, obj, natoms, coords, atype, box, suffix): t_out, feed_dict = build_tf_descriptor(obj, natoms, coords, atype, box, suffix) t_out_indentity = [ tf.identity(tt, name=f"o_{ii}_{suffix}") for ii, tt in enumerate(t_out) ] - run_sess(sess, tf.global_variables_initializer()) - return run_sess( - sess, - t_out_indentity, - feed_dict=feed_dict, - ) + try: + run_sess(sess, tf.global_variables_initializer()) + return run_sess( + sess, + t_out_indentity, + feed_dict=feed_dict, + ) + except tf.errors.OpError as e: + raise RuntimeError(f"Failed to evaluate descriptor: {str(e)}")
62-116: Remove or document the commented-out test case.There's a commented-out alternative
atypearray that might indicate an incomplete test case:# self.atype = np.array([0, 0, 1, 1, 1, 1], dtype=np.int32)Either remove it or document why it's kept for future reference.
deepmd/tf/descriptor/se_atten.py (3)
2262-2262: Fix typos in the docstring.The docstring contains typos: 'Reveive' should be 'Receive', and 'statisitcs' should be 'statistics'.
2283-2283: Correct typo in the comment.In the comment, 'mocel' should be 'model'.
2251-2348: Add unit tests forenable_compressioninDescrptDPA1Compat.Ensure that the new
enable_compressionmethod inDescrptDPA1Compatis covered by unit tests to verify its functionality and catch potential issues.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
- deepmd/tf/descriptor/se_atten.py (1 hunks)
- deepmd/tf/utils/compress.py (1 hunks)
- deepmd/tf/utils/tabulate.py (1 hunks)
- source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py (1 hunks)
🧰 Additional context used
🔇 Additional comments (5)
deepmd/tf/utils/compress.py (1)
23-24: LGTM! The suffix parameter addition looks good.The change maintains backward compatibility while enabling flexible tensor naming needed for compression support.
Let's verify the usage of this method with different suffixes:
✅ Verification successful
Based on the shell script results, I can now generate the final response:
The suffix parameter is correctly used and integrated across the codebase
The verification shows that:
- In
se_atten.py, the method is called with an explicit suffix parameter:get_two_side_type_embedding(self, graph, suffix=suffix)- In
se_a.py, it's called without a suffix, utilizing the default empty string- The tensor naming pattern
t_typeebd{suffix}is consistently used across the codebase:
- In
model.py:f"t_typeebd{suffix}:0"- In
type_embed.py:name="t_typeebd" + suffix- In test files, the base name
t_typeebdis used without suffix🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check how the get_two_side_type_embedding method is called across the codebase # and verify that suffix parameter is used consistently # Search for method calls rg -A 3 "get_two_side_type_embedding\(" # Search for tensor name patterns to ensure consistent suffix usage rg "t_typeebd[^\"]*"Length of output: 5424
source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py (3)
1-21: LGTM: Imports are well-organized and complete.The imports are properly structured and include all necessary dependencies for the test implementation.
23-45: LGTM: Well-structured descriptor building function.The function correctly sets up TensorFlow placeholders with appropriate types and names, and properly handles the suffix parameter.
153-154: LGTM: Standard unittest main block.The main block is correctly implemented for running the tests.
deepmd/tf/utils/tabulate.py (1)
129-131: LGTM! Consistent tensor name resolution.The change correctly adds the suffix to the tensor name when retrieving
ntypes, maintaining consistency with how other tensors (davg,dstd) are retrieved from the graph.
source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py
Outdated
Show resolved
Hide resolved
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
deepmd/tf/descriptor/se_atten.py (1)
Line range hint
426-449: LGTM! Consider enhancing parameter documentation.The addition of
tebd_suffixparameter provides flexibility in scope naming for type embedding during model compression. The parameter documentation could be more specific about when to use a different suffix than the main scope suffix.Consider updating the docstring to clarify:
@@ -2288,7 +2288 @@ tebd_suffix : str, optional - The suffix of the type embedding scope, only for DescrptDPA1Compat + The suffix of the type embedding scope. When using DescrptDPA1Compat, + this must be the same as the main suffix. For other implementations, + it can be different to support separate scoping of type embeddings.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
- deepmd/tf/descriptor/se_atten.py (4 hunks)
🧰 Additional context used
🔇 Additional comments (2)
deepmd/tf/descriptor/se_atten.py (2)
502-504: LGTM! Consistent suffix propagation.The
tebd_suffixis correctly passed toget_two_side_type_embedding, maintaining consistency with the method signature change.
2256-2305: LGTM! Well-constrained implementation.The override correctly enforces that
tebd_suffixmust matchsuffixinDescrptDPA1Compat, which is essential for maintaining compatibility. The parent method is called appropriately with the same suffix.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
deepmd/tf/descriptor/se_atten.py (1)
2257-2306: Enhance docstring to explain suffix constraintThe implementation correctly enforces that
tebd_suffixmust matchsuffixinDescrptDPA1Compat. However, the docstring could be improved to explain why this constraint exists.Consider updating the docstring for
tebd_suffixto:- Same as suffix. + The suffix of the type embedding scope. Must be the same as `suffix` in DescrptDPA1Compat + to maintain compatibility with the original implementation.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
- deepmd/tf/descriptor/se_atten.py (5 hunks)
- source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py
🧰 Additional context used
🪛 Ruff
deepmd/tf/descriptor/se_atten.py
567-567: Local variable
t_ntypesis assigned to but never usedRemove assignment to unused variable
t_ntypes(F841)
🔇 Additional comments (2)
deepmd/tf/descriptor/se_atten.py (2)
Line range hint
426-449: LGTM: Well-documented parameter additionThe addition of the
tebd_suffixparameter toenable_compressionis well-implemented and properly documented. This change provides flexibility in specifying the suffix for type embedding scope.Also applies to: 502-504
566-567: Verify the purpose of unused variablet_ntypesThe variable
t_ntypesis assigned but never used. However, since it's in a TensorFlow graph context, it might serve a purpose in the graph structure or debugging.Let's check if this variable is used in the graph:
🧰 Tools
🪛 Ruff
567-567: Local variable
t_ntypesis assigned to but never usedRemove assignment to unused variable
t_ntypes(F841)
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
deepmd/tf/descriptor/se_atten.py (1)
567-568: Keep unused variables for documentation and debugging.While
t_ntypesandt_rcutare currently unused, they serve as important metadata about the model's configuration and may be useful for debugging. Consider adding a comment explaining their purpose or marking them with# pylint: disable=unused-variableif they are intentionally kept.🧰 Tools
🪛 Ruff
567-567: Local variable
t_ntypesis assigned to but never usedRemove assignment to unused variable
t_ntypes(F841)
568-568: Local variable
t_rcutis assigned to but never usedRemove assignment to unused variable
t_rcut(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
- deepmd/tf/descriptor/se_atten.py (5 hunks)
🧰 Additional context used
🪛 Ruff
deepmd/tf/descriptor/se_atten.py
567-567: Local variable
t_ntypesis assigned to but never usedRemove assignment to unused variable
t_ntypes(F841)
568-568: Local variable
t_rcutis assigned to but never usedRemove assignment to unused variable
t_rcut(F841)
🔇 Additional comments (2)
deepmd/tf/descriptor/se_atten.py (2)
Line range hint
426-504: LGTM: Proper handling of type embedding suffix in compression.The addition of the
tebd_suffixparameter and its usage inget_two_side_type_embeddingcorrectly addresses the issue with compression suffix handling in type embeddings.
2257-2306: LGTM: Enforcing consistent suffix usage in DescrptDPA1Compat.The implementation correctly enforces that
tebd_suffixmust matchsuffixin DescrptDPA1Compat, which is essential for maintaining consistency in the compression process.
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## devel #4243 +/- ##
==========================================
- Coverage 84.58% 84.28% -0.30%
==========================================
Files 547 548 +1
Lines 51327 51434 +107
Branches 3047 3051 +4
==========================================
- Hits 43413 43353 -60
- Misses 6967 7120 +153
- Partials 947 961 +14 ☔ View full report in Codecov by Sentry. |
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
Fix #4114 .
Summary by CodeRabbit
New Features
Bug Fixes
Tests