fix: use tuple in xp.reshape
#4808
Conversation
📝 WalkthroughWalkthroughThe update standardizes array shape specifications by replacing list-based shape arguments with tuple-based ones in calls to Changes
Sequence Diagram(s)No sequence diagram generated as the changes are limited to array shape syntax and do not affect control flow or feature logic. Suggested reviewers
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Pylint (3.3.7)deepmd/dpmodel/utils/env_mat_stat.pyNo files to lint: exiting. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (29)
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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 (
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
deepmd/dpmodel/fitting/general_fitting.py (4)
415-416: Considerbroadcast_toto avoid thetilememory-duplication
xp.tile(xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1))materialises a full(nf, nloc, numb_fparam)array.
With the same semantics you can keep everything as a lightweight broadcasted view:-fparam = xp.tile( - xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1) -) +fparam = xp.broadcast_to( + xp.reshape(fparam, (nf, 1, self.numb_fparam)), + (nf, nloc, self.numb_fparam), +)Reduces both peak memory and unnecessary copies.
449-450: Sametilecaveat as above forcase_embdcase_embd = xp.tile( xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1) )
xp.broadcast_towould achieve the same broadcast semantics without allocating the repeated data:-case_embd = xp.tile( - xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1) -) +case_embd = xp.broadcast_to( + xp.reshape(self.case_embd[...], (1, 1, -1)), + (nf, nloc, self.dim_case_embd), +)Minor, but can matter for large systems.
468-469: Mask construction: use boolean broadcasting instead oftileThe current pattern duplicates a
(nf, nloc)boolean masknet_dim_outtimes:mask = xp.tile(xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out))This can be simplified and zero-copy by relying on implicit broadcasting:
-mask = xp.tile( - xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out) -) +mask = xp.reshape((atype == type_i), (nf, nloc, 1))All later
xp.where(mask, …)calls will automatically broadcast the singleton dimension acrossnet_dim_out.
488-492: Type-safety check foratypebeforexp.take
xp.takerequiresatypeto be an integer tensor. If callers accidentally pass a float array the code will silently cast on some back-ends but raise on others. Consider enforcing:if not xp.issubdtype(atype.dtype, xp.integer): raise TypeError("`atype` must contain integer type indices")Not a blocker, but will make failures deterministic across back-ends.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/fitting/general_fitting.py(5 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
🔇 Additional comments (1)
deepmd/dpmodel/fitting/general_fitting.py (1)
434-435: LGTM – shape clarificationSwitching to a tuple literal in
xp.reshape(aparam, (nf, nloc, self.numb_aparam))keeps the behaviour intact and aligns with the Array-API spec.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
deepmd/dpmodel/utils/exclude_mask.py (1)
55-57: Align reshape-API usage across the classYou correctly switched the runtime reshape to use a tuple, but the class-level cache created in
__init__still uses the old list style (reshape([-1]), line 26). Small, but it leaves an inconsistency in the public attributetype_mask.- self.type_mask = type_mask.reshape([-1]) + # Keep reshape syntax consistent with the rest of the codebase + self.type_mask = type_mask.reshape((-1,))deepmd/dpmodel/utils/region.py (1)
95-97: Prefer unpacking over list-conversion for clearer intent
tuple(list(cshape[:-2]) + [3])allocates an intermediate list every call. You can achieve the same shape without the extra object and with simpler intent:- return xp.reshape(dist, tuple(list(cshape[:-2]) + [3])) # noqa:RUF005 + return xp.reshape(dist, (*cshape[:-2], 3))Besides the tiny performance gain, using the star-unpack pattern is more idiomatic and easier to read.
deepmd/dpmodel/fitting/polarizability_fitting.py (1)
314-322: Minor API-compatibility concern regardingxp.astypeWhile untouched by this diff, note that
array_api_compatnamespaces do not guarantee a top-levelastypefunction (NumPy providesnp.asarray(..., dtype=...)instead). If any backend other than NumPy is used, this may raise. Consider switching toxp.asarray(self.constant_matrix, dtype=out.dtype)for maximum array-API compliance.
deepmd/dpmodel/loss/ener.py (1)
239-243: Minor style consistencyTo stay uniform with the rest of the edits, you may also convert the shape in this block to tuples:
- drdq_reshape = xp.reshape( - drdq, (-1, natoms[0] * 3, self.numb_generalized_coord) - ) + drdq_reshape = xp.reshape( + drdq, (-1, natoms[0] * 3, self.numb_generalized_coord) + )(Not critical—just style.)
deepmd/dpmodel/model/transform_output.py (1)
98-112: Consider avoiding shadowing ofmappingfor clarityInside the
for-loop,mappingis reassigned twice (reshapethentile), while the original array is kept asmapping_. Although functionally correct, this shadowing can be confusing and error-prone in future edits.- mapping = xp.reshape( - mapping, tuple(mldims + [1] * len(derv_r_ext_dims)) - ) + mapping_reshaped = xp.reshape( + mapping, tuple(mldims + [1] * len(derv_r_ext_dims)) + ) ... - mapping = xp.tile(mapping, [1] * len(mldims) + derv_r_ext_dims) + mapping_tiled = xp.tile( + mapping_reshaped, [1] * len(mldims) + derv_r_ext_dims + )Renaming the intermediate tensors makes the intent explicit and prevents accidental misuse later in the block (and in the
c_differentiablesection).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
deepmd/dpmodel/descriptor/dpa1.py(1 hunks)deepmd/dpmodel/descriptor/dpa2.py(1 hunks)deepmd/dpmodel/descriptor/dpa3.py(1 hunks)deepmd/dpmodel/descriptor/se_t_tebd.py(1 hunks)deepmd/dpmodel/fitting/polarizability_fitting.py(2 hunks)deepmd/dpmodel/loss/ener.py(5 hunks)deepmd/dpmodel/model/transform_output.py(1 hunks)deepmd/dpmodel/utils/env_mat_stat.py(1 hunks)deepmd/dpmodel/utils/exclude_mask.py(1 hunks)deepmd/dpmodel/utils/neighbor_stat.py(1 hunks)deepmd/dpmodel/utils/nlist.py(1 hunks)deepmd/dpmodel/utils/region.py(1 hunks)
✅ Files skipped from review due to trivial changes (6)
- deepmd/dpmodel/utils/env_mat_stat.py
- deepmd/dpmodel/descriptor/dpa2.py
- deepmd/dpmodel/descriptor/dpa1.py
- deepmd/dpmodel/utils/neighbor_stat.py
- deepmd/dpmodel/descriptor/se_t_tebd.py
- deepmd/dpmodel/utils/nlist.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
deepmd/dpmodel/loss/ener.py (1)
deepmd/dpmodel/loss/loss.py (1)
display_if_exist(38-56)
⏰ Context from checks skipped due to timeout of 90000ms (24)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (6)
deepmd/dpmodel/fitting/polarizability_fitting.py (1)
290-294: Change looks goodTuple reshape matches the PR objective and the rest of the codebase. No further issues spotted.
deepmd/dpmodel/loss/ener.py (2)
135-147:xp.normis not part of the Array-APIThe new tuple reshapes are fine, but line 143 still calls
xp.norm, which is outside the official array-API spec and will fail for backends that expose only the standard namespace. Usexp.linalg.vector_norm(orxp.linalg.norm) instead.- norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), (-1, 1)) + self.relative_f + norm_f = xp.reshape( + xp.linalg.vector_norm(force_hat_3, axis=1), (-1, 1) + ) + self.relative_f
186-189: Tuple reshape LGTMThe huber-loss branch now uses tuple style consistently—good catch.
deepmd/dpmodel/descriptor/dpa3.py (2)
563-567: Reshape update approvedTuple reshape is correct and keeps the backend-agnostic indexing intact. No further issues.
569-572: Same as aboveConsistent tuple usage—looks good.
deepmd/dpmodel/model/transform_output.py (1)
103-105: Correct use of tuple forxp.reshape👍Switching the second argument of
xp.reshapefrom a bare list to an explicittuple(...)fixes compatibility issues with back-ends (e.g. some JAX / CuPy versions reject list shapes). No further action needed.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/dpmodel/array_api.py (1)
61-64: Build the target shape without the intermediate list allocation
arr.shapeis already a tuple; turning it into a list, mutating, and then converting back introduces an avoidable copy and an extra mutation step.- shape = list(arr.shape) - shape.pop(-1) - shape = (*shape, n) + # keep it tuple throughout + shape = (*arr.shape[:-1], n)This single-expression version is clearer, avoids temporary objects, and remains backend-agnostic.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/array_api.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
There was a problem hiding this comment.
Pull Request Overview
This PR standardizes array shape specifications in reshape and take operations by using tuples instead of lists, improving style consistency without altering functionality.
- Consistently replace list-based shape arguments in
xp.reshapeandxp.takewith tuple-based arguments. - Minor refactoring in
deepmd/dpmodel/array_api.pyto use tuple unpacking for shape construction.
Reviewed Changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| deepmd/dpmodel/utils/region.py | Use tuple for reshape shapes |
| deepmd/dpmodel/utils/nlist.py | Use tuple for reshape shapes |
| deepmd/dpmodel/utils/neighbor_stat.py | Use tuple for reshape shapes |
| deepmd/dpmodel/utils/exclude_mask.py | Use tuple for reshape shapes in xp.take |
| deepmd/dpmodel/utils/env_mat_stat.py | Use tuple for reshape shapes |
| deepmd/dpmodel/model/transform_output.py | Use tuple for multi-dimensional reshape |
| deepmd/dpmodel/loss/ener.py | Use tuple for reshape shapes throughout loss logic |
| deepmd/dpmodel/fitting/polarizability_fitting.py | Use tuple for reshape in scaling operations |
| deepmd/dpmodel/fitting/general_fitting.py | Use tuple for reshape in parameter tiling |
| deepmd/dpmodel/descriptor/se_t_tebd.py | Use tuple for reshape in type embedding expansion |
| deepmd/dpmodel/descriptor/dpa3.py | Use tuple for reshape in node embedding expansion |
| deepmd/dpmodel/descriptor/dpa2.py | Use tuple for reshape in generic embedding |
| deepmd/dpmodel/descriptor/dpa1.py | Use tuple for reshape in atomic embedding |
| deepmd/dpmodel/array_api.py | Use tuple unpacking for dynamic shape construction |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4808 +/- ##
==========================================
- Coverage 84.80% 84.79% -0.01%
==========================================
Files 698 698
Lines 67798 67816 +18
Branches 3542 3541 -1
==========================================
+ Hits 57494 57503 +9
- Misses 9171 9177 +6
- Partials 1133 1136 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
1dc5b04
This PR fixes the error in UT where the argument of `xp.reshape` should be a tuple, not a list. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Style** - Standardized the use of tuple syntax for shape arguments in array reshaping operations throughout the application, replacing previous list-based syntax. This change ensures consistency and aligns with best practices for array manipulation. No functional behavior is affected. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This PR fixes the error in UT where the argument of
xp.reshapeshould be a tuple, not a list.Summary by CodeRabbit