refactor(dpmodel/dpa4)#5599
Open
OutisLi wants to merge 5 commits into
Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
📝 WalkthroughWalkthroughAdds Cartesian tensor-product routing for SeZM/DPA4, refactors dpmodel building blocks and serialization, introduces LoRA adapters, adds pt_expt Triton/runtime overrides, and updates compile and test coverage for the new descriptor paths. ChangesSeZM/DPA4 descriptor extension and dpmodel port refactor
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/dpmodel/descriptor/dpa4_nn/norm.py (1)
354-361: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRestore
mmaxbounds validation before deriving balance weights.With
mmax < 0,n_coeff_lcan become non-positive, producing invalid balance weights and NaNs incall(). Keep the constructor aligned with the documented0 <= mmax <= lmaxcontract.Proposed fix
self.lmax = int(lmax) self.mmax = int(mmax) + if self.lmax < 0: + raise ValueError("`lmax` must be non-negative") + if self.mmax < 0: + raise ValueError("`mmax` must be non-negative") + if self.mmax > self.lmax: + raise ValueError("`mmax` must be <= `lmax`") self.channels = int(channels)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/descriptor/dpa4_nn/norm.py` around lines 354 - 361, Restore the missing `mmax` validation in the `Norm.__init__` constructor before any balance-weight derivation happens. Ensure the constructor still enforces the documented `0 <= mmax <= lmax` contract, alongside the existing `lmax` check, so `n_coeff_l` cannot become non-positive and propagate invalid values into `call()`. Use the `Norm` initializer and the `mmax`/`lmax` guard logic as the place to fix this.deepmd/dpmodel/descriptor/dpa4_nn/attention.py (1)
77-78: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRestore the padded-layout shape guard before reshaping.
nnei = n_edge // n_nodesnow turnsn_nodes == 0intoZeroDivisionError, and non-multipleEvalues fail later during reshape. Keep the explicit contract check here so callers get the intended error.Proposed fix
n_channel = n_focus * n_head eps_f = float(eps) + if n_nodes <= 0 or n_edge % n_nodes != 0: + raise ValueError("`logits.shape[0]` must be a positive multiple of `n_nodes`") nnei = n_edge // n_nodes🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/descriptor/dpa4_nn/attention.py` around lines 77 - 78, Restore the explicit padded-layout shape validation in attention.py before computing nnei or reshaping: the logic around eps_f, n_edge, and n_nodes should reject invalid inputs up front instead of relying on n_edge // n_nodes. Add back the contract check in the attention reshaping path so n_nodes == 0 and non-multiple edge counts raise the intended error, and keep this guard near the code that derives nnei and reshapes the attention tensors.
🧹 Nitpick comments (2)
deepmd/dpmodel/descriptor/dpa4_nn/so3.py (1)
505-505: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winValidate deterministic
expand_indexon load.
expand_indexis fully determined bylmax; accepting the serialized value directly can silently corrupt degree-to-weight mapping if a checkpoint is stale or malformed but shape-compatible.Proposed fix
- obj.expand_index = np.asarray(variables["expand_index"], dtype=np.int64) + loaded_expand_index = np.asarray(variables["expand_index"], dtype=np.int64) + expected_expand_index = map_degree_idx(obj.lmax) + if not np.array_equal(loaded_expand_index, expected_expand_index): + raise ValueError("Invalid SO3Linear expand_index for configured lmax") + obj.expand_index = expected_expand_index🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/descriptor/dpa4_nn/so3.py` at line 505, The load path for `SO3` is trusting the serialized `expand_index` value directly, even though it is fully derived from `lmax`. Update the deserialization logic in `so3.py` so `expand_index` is recomputed from the model’s `lmax` during load (or at least validated against the expected deterministic value) instead of accepting the checkpoint contents blindly. Keep the change localized to the object restore flow around `obj.expand_index` to prevent stale or malformed checkpoints from overriding the correct degree-to-weight mapping.deepmd/pt_expt/descriptor/dpa4_nn/radial.py (1)
73-78: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFail loudly on unexpected
RadialMLP.netentries.This fallback currently turns every non-
Module/non-NativeOPentry into_ScalarActivation. If the dpmodel list shape changes later, this wrapper will silently build the wrong module tree.Possible tightening
def _convert_layer(self, layer: Any) -> torch.nn.Module: if isinstance(layer, torch.nn.Module): return layer if isinstance(layer, NativeOP): return try_convert_module(layer) - return _ScalarActivation(self.activation_function) + if callable(layer): + return _ScalarActivation(self.activation_function) + raise TypeError(f"Unsupported RadialMLP.net entry: {type(layer)!r}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt_expt/descriptor/dpa4_nn/radial.py` around lines 73 - 78, The _convert_layer method in RadialMLP currently falls back to _ScalarActivation for any unexpected entry, which can silently mask bad dpmodel shapes. Tighten the logic so only the known cases (torch.nn.Module and NativeOP) are accepted and raise an explicit error for anything else, using _convert_layer and RadialMLP.net as the key locations to update.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/descriptor/dpa4_nn/attn_res.py`:
- Around line 163-176: The main issue is that the `forward` logic in
`attn_res.py` dereferences `sources[0]` and calls `scalar_extractor(current_x)`
without validating required inputs first. Add explicit guards near the start of
the method to reject an empty `sources` sequence and to require `current_x`
whenever `self.input_dependent` is enabled, using the existing
`scalar_extractor` and `query_proj` path as the location to enforce the
contract. Raise a clear, descriptive error before any indexing or projection
happens so invalid calls fail fast instead of producing opaque runtime errors.
In `@deepmd/dpmodel/descriptor/dpa4_nn/block.py`:
- Around line 297-302: The dpmodel deserialization path is not
backward-compatible because older configs still use so2_layers while the
constructor now expects mixing_layers. Update the block’s config-loading logic
in the relevant deserialization/compatibility entry point for DPA4 block classes
(including the paths around the referenced constructor and any shared loaders)
to translate legacy so2_layers into mixing_layers before calling cls(**config).
Keep the old key readable, preserve existing behavior when mixing_layers is
already present, and avoid passing both keys into the constructor.
In `@deepmd/dpmodel/descriptor/dpa4_nn/embedding.py`:
- Around line 257-258: The reshape-based aggregation in embedding.py is only
valid for the padded build_edge_cache layout, so the sparse
build_edge_cache_from_edges path must not use `nnei = n_edge // n_nodes` or the
related reshape logic. Update the relevant aggregation methods in the dpa4_nn
embedding flow to either keep the scatter/index-add implementation for sparse
edge caches or explicitly guard these reshape paths so they only run on padded
inputs. Use the existing edge-cache handling in the descriptor/embedding methods
to distinguish padded versus sparse cases and apply the same rule to the other
reshape-based aggregations in this file.
In `@deepmd/dpmodel/descriptor/dpa4_nn/lora.py`:
- Around line 652-660: The tuple-handling in _swap_submodule is incorrect
because it treats tuple children like mutable lists and tries in-place
replacement. Update the replacement logic in _swap_submodule to either only
allow indexed assignment for list parents or, when the parent is a tuple,
rebuild the tuple with the new NativeOP and reassign it back to the container
instead of mutating it directly.
- Around line 706-739: The LoRA adapter parameters are ending up frozen because
apply_lora_to_sezm freezes every module first and the replacement
LoRASO3/LoRASO2 instances inherit that state from their serialized config.
Update the replacement logic in apply_lora_to_sezm so the injected LoRASO3 and
LoRASO2 adapters are explicitly marked trainable after construction, while
keeping the copied base weights/biases frozen under the dpmodel trainability
model. Make sure the same fix is applied in the later replacement block that
handles the other SO3Linear/SO2Linear occurrences.
In `@deepmd/dpmodel/descriptor/dpa4_nn/projection.py`:
- Around line 278-305: The e3nn quadrature setup in the projection path assumes
an even theta resolution, but `_normalize_s2_grid_resolution()` can still pass
through a custom odd `res_beta`, causing `quad_weight` to have the wrong length
and `harmonics` broadcasting to fail. Update the resolution handling in this
projection code to either reject odd e3nn `theta_resolution` values with a clear
error or normalize them to the next even value before computing `quad_weight`,
`to_grid_mat`, and `from_grid_mat`.
In `@deepmd/pt_expt/descriptor/dpa4_nn/block.py`:
- Around line 32-34: The pt_expt descriptor block is importing the dpmodel-only
exchange_ghost_features helper, which is not implemented for torch and will fail
whenever comm_dict is used. Update the import in block.py to use the
PyTorch-capable ghost exchange implementation, or add a pt_expt wrapper that
delegates to the torch version, and make sure the callers in the SO(2) message
passing path use that symbol instead of the dpmodel one.
In `@deepmd/pt_expt/descriptor/dpa4_nn/triton/radial_mix.py`:
- Around line 687-698: The backward launch in _radial_mix_bwd_kernel is mixing a
contiguous grad buffer with strides taken from the original grad_out, which can
produce incorrect reads when grad_out is non-contiguous. Update the Triton call
so the stride arguments come from the same contiguous tensor being passed into
the kernel, using the grad_out.contiguous() value consistently for the stride
calculations.
In `@deepmd/pt/model/descriptor/sezm_nn/block.py`:
- Around line 303-308: The block deserialization path is not backward compatible
because legacy serialized configs may still use so2_layers while the constructor
for the block class now expects mixing_layers. Update the config migration in
the block loading/deserialization flow (around the block constructor/init path
and any versioned load logic) to translate so2_layers into mixing_layers before
calling cls(**config), or explicitly branch on the serialized version and map
old payloads in the same block class. Ensure this compatibility handling is
applied in both affected load locations.
---
Outside diff comments:
In `@deepmd/dpmodel/descriptor/dpa4_nn/attention.py`:
- Around line 77-78: Restore the explicit padded-layout shape validation in
attention.py before computing nnei or reshaping: the logic around eps_f, n_edge,
and n_nodes should reject invalid inputs up front instead of relying on n_edge
// n_nodes. Add back the contract check in the attention reshaping path so
n_nodes == 0 and non-multiple edge counts raise the intended error, and keep
this guard near the code that derives nnei and reshapes the attention tensors.
In `@deepmd/dpmodel/descriptor/dpa4_nn/norm.py`:
- Around line 354-361: Restore the missing `mmax` validation in the
`Norm.__init__` constructor before any balance-weight derivation happens. Ensure
the constructor still enforces the documented `0 <= mmax <= lmax` contract,
alongside the existing `lmax` check, so `n_coeff_l` cannot become non-positive
and propagate invalid values into `call()`. Use the `Norm` initializer and the
`mmax`/`lmax` guard logic as the place to fix this.
---
Nitpick comments:
In `@deepmd/dpmodel/descriptor/dpa4_nn/so3.py`:
- Line 505: The load path for `SO3` is trusting the serialized `expand_index`
value directly, even though it is fully derived from `lmax`. Update the
deserialization logic in `so3.py` so `expand_index` is recomputed from the
model’s `lmax` during load (or at least validated against the expected
deterministic value) instead of accepting the checkpoint contents blindly. Keep
the change localized to the object restore flow around `obj.expand_index` to
prevent stale or malformed checkpoints from overriding the correct
degree-to-weight mapping.
In `@deepmd/pt_expt/descriptor/dpa4_nn/radial.py`:
- Around line 73-78: The _convert_layer method in RadialMLP currently falls back
to _ScalarActivation for any unexpected entry, which can silently mask bad
dpmodel shapes. Tighten the logic so only the known cases (torch.nn.Module and
NativeOP) are accepted and raise an explicit error for anything else, using
_convert_layer and RadialMLP.net as the key locations to update.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5f98d8e6-9f05-462e-abee-46098a57be4e
📒 Files selected for processing (52)
deepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/__init__.pydeepmd/dpmodel/descriptor/dpa4_nn/activation.pydeepmd/dpmodel/descriptor/dpa4_nn/attention.pydeepmd/dpmodel/descriptor/dpa4_nn/attn_res.pydeepmd/dpmodel/descriptor/dpa4_nn/block.pydeepmd/dpmodel/descriptor/dpa4_nn/cartesian.pydeepmd/dpmodel/descriptor/dpa4_nn/edge_cache.pydeepmd/dpmodel/descriptor/dpa4_nn/embedding.pydeepmd/dpmodel/descriptor/dpa4_nn/ffn.pydeepmd/dpmodel/descriptor/dpa4_nn/grid_net.pydeepmd/dpmodel/descriptor/dpa4_nn/indexing.pydeepmd/dpmodel/descriptor/dpa4_nn/lora.pydeepmd/dpmodel/descriptor/dpa4_nn/norm.pydeepmd/dpmodel/descriptor/dpa4_nn/projection.pydeepmd/dpmodel/descriptor/dpa4_nn/radial.pydeepmd/dpmodel/descriptor/dpa4_nn/so2.pydeepmd/dpmodel/descriptor/dpa4_nn/so3.pydeepmd/dpmodel/descriptor/dpa4_nn/utils.pydeepmd/dpmodel/descriptor/dpa4_nn/wignerd.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/__init__.pydeepmd/pt/model/descriptor/sezm_nn/block.pydeepmd/pt/model/descriptor/sezm_nn/cartesian.pydeepmd/pt/model/descriptor/sezm_nn/edge_cache.pydeepmd/pt/model/descriptor/sezm_nn/so2.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/utils/compile_compat.pydeepmd/pt_expt/descriptor/dpa4.pydeepmd/pt_expt/descriptor/dpa4_nn/__init__.pydeepmd/pt_expt/descriptor/dpa4_nn/block.pydeepmd/pt_expt/descriptor/dpa4_nn/radial.pydeepmd/pt_expt/descriptor/dpa4_nn/so2.pydeepmd/pt_expt/descriptor/dpa4_nn/triton/__init__.pydeepmd/pt_expt/descriptor/dpa4_nn/triton/radial_mix.pydeepmd/pt_expt/descriptor/dpa4_nn/triton/so2_rotation.pydeepmd/utils/argcheck.pyexamples/water/dpa4/input.jsonsource/tests/common/dpmodel/test_descrpt_dpa4.pysource/tests/common/dpmodel/test_dpa4_basegridnet_cross.pysource/tests/common/dpmodel/test_dpa4_frame_mixers.pysource/tests/common/dpmodel/test_dpa4_grid_descriptor.pysource/tests/common/dpmodel/test_dpa4_gridmlp_frames.pysource/tests/common/dpmodel/test_dpa4_so2_grid.pysource/tests/common/dpmodel/test_dpa4_so3_grid_utils.pysource/tests/common/dpmodel/test_dpa4_so3_gridnet.pysource/tests/common/dpmodel/test_lebedev_sh.pysource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_dpa4_dpmodel_parity.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt_expt/descriptor/test_dpa4.pysource/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py
💤 Files with no reviewable changes (4)
- source/tests/common/dpmodel/test_dpa4_frame_mixers.py
- source/tests/pt_expt/descriptor/test_dpa4.py
- source/tests/common/dpmodel/test_dpa4_basegridnet_cross.py
- source/tests/common/dpmodel/test_lebedev_sh.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5599 +/- ##
==========================================
- Coverage 82.35% 81.62% -0.73%
==========================================
Files 896 908 +12
Lines 100952 103436 +2484
Branches 4059 4058 -1
==========================================
+ Hits 83138 84434 +1296
- Misses 16349 17540 +1191
+ Partials 1465 1462 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/descriptor/dpa4_nn/embedding.py (1)
147-156: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winNormalize
precisionbefore serializing.These serializers index
PRECISION_DICTwithself.precision, but the constructors accept mixed-case values viaself.precision.lower().SeZMTypeEmbedding(precision="Float64").serialize()now raisesKeyError, and the same regression repeats inGeometricInitialEmbeddingandEnvironmentInitialEmbedding.Proposed fix
- "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "precision": np.dtype(PRECISION_DICT[self.precision.lower()]).name,- "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "precision": np.dtype(PRECISION_DICT[self.precision.lower()]).name,- "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "precision": np.dtype(PRECISION_DICT[self.precision.lower()]).name,Also applies to: 336-343, 667-683
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/descriptor/dpa4_nn/embedding.py` around lines 147 - 156, The serializer methods for SeZMTypeEmbedding, GeometricInitialEmbedding, and EnvironmentInitialEmbedding are indexing PRECISION_DICT with the raw precision value, which can fail when callers pass mixed-case strings that were accepted by the constructors. Update each serialize() implementation to normalize self.precision before the PRECISION_DICT lookup, matching the lowercasing already used in the corresponding constructors, so serialization works consistently for values like "Float64".
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@deepmd/dpmodel/descriptor/dpa4_nn/embedding.py`:
- Around line 147-156: The serializer methods for SeZMTypeEmbedding,
GeometricInitialEmbedding, and EnvironmentInitialEmbedding are indexing
PRECISION_DICT with the raw precision value, which can fail when callers pass
mixed-case strings that were accepted by the constructors. Update each
serialize() implementation to normalize self.precision before the PRECISION_DICT
lookup, matching the lowercasing already used in the corresponding constructors,
so serialization works consistently for values like "Float64".
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 9cdaa53a-d8f8-4924-84b5-e04c7fc2562d
📒 Files selected for processing (11)
deepmd/dpmodel/array_api.pydeepmd/dpmodel/descriptor/dpa4_nn/attention.pydeepmd/dpmodel/descriptor/dpa4_nn/embedding.pydeepmd/dpmodel/descriptor/dpa4_nn/ffn.pydeepmd/dpmodel/descriptor/dpa4_nn/lora.pydeepmd/dpmodel/descriptor/dpa4_nn/so2.pydeepmd/pt_expt/descriptor/__init__.pydeepmd/pt_expt/descriptor/dpa4.pysource/tests/common/dpmodel/test_dpa4_lora.pysource/tests/common/dpmodel/test_dpa4_sparse_edges.pysource/tests/pt/model/test_dpa4_dpmodel_parity.py
💤 Files with no reviewable changes (2)
- deepmd/pt_expt/descriptor/dpa4.py
- deepmd/dpmodel/descriptor/dpa4_nn/ffn.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/descriptor/dpa4_nn/lora.py
- source/tests/pt/model/test_dpa4_dpmodel_parity.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt_expt/utils/serialization.py (1)
1009-1045: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the inductor config save/disable/restore into a context manager.
The threshold +
assert_indirect_indexingsnapshot/override/restore block is now duplicated verbatim in both the regular and with-comm compile paths. A small context manager centralizes the documented thread-safety caveat in one place and prevents the two copies from silently diverging later.♻️ Proposed context-manager extraction
from contextlib import contextmanager `@contextmanager` def _inductor_compile_overrides(*, disable_fusion: bool): # See module notes: torch._inductor.config is a process-wide singleton; # this save/restore is NOT thread-safe. Serialise .pt2 exports per process. import torch._inductor.config as _inductor_config saved_threshold = _inductor_config.realize_opcount_threshold saved_assert_indexing = _inductor_config.assert_indirect_indexing if disable_fusion: _inductor_config.realize_opcount_threshold = 0 _inductor_config.assert_indirect_indexing = False try: yield finally: _inductor_config.realize_opcount_threshold = saved_threshold _inductor_config.assert_indirect_indexing = saved_assert_indexingUsage at both call sites:
- saved_threshold = _inductor_config.realize_opcount_threshold - saved_assert_indexing = _inductor_config.assert_indirect_indexing - if is_cuda: - _inductor_config.realize_opcount_threshold = 0 - _inductor_config.assert_indirect_indexing = False - try: - aoti_compile_and_package(exported, package_path=model_file) - finally: - _inductor_config.realize_opcount_threshold = saved_threshold - _inductor_config.assert_indirect_indexing = saved_assert_indexing + with _inductor_compile_overrides(disable_fusion=is_cuda): + aoti_compile_and_package(exported, package_path=model_file)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt_expt/utils/serialization.py` around lines 1009 - 1045, The inductor config save/override/restore logic is duplicated in both `aoti_compile_and_package` compile paths, so extract it into a shared context manager (for example, a helper near the serialization/export code) to keep the process-wide `torch._inductor.config` handling in one place. Move the `realize_opcount_threshold` and `assert_indirect_indexing` snapshot/restore behavior into that helper, then wrap both the regular export and the `with_comm` export call sites with it so the behavior stays identical and cannot diverge.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@deepmd/pt_expt/utils/serialization.py`:
- Around line 1009-1045: The inductor config save/override/restore logic is
duplicated in both `aoti_compile_and_package` compile paths, so extract it into
a shared context manager (for example, a helper near the serialization/export
code) to keep the process-wide `torch._inductor.config` handling in one place.
Move the `realize_opcount_threshold` and `assert_indirect_indexing`
snapshot/restore behavior into that helper, then wrap both the regular export
and the `with_comm` export call sites with it so the behavior stays identical
and cannot diverge.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 92be69fa-51ce-4796-821c-790f683b1700
📒 Files selected for processing (1)
deepmd/pt_expt/utils/serialization.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is first PR of porting the whole dpa4 descriptor
Summary by CodeRabbit
New Features
edge_cartesian,node_cartesian) withmixing_layerscontrol.Bug Fixes
API Changes
so2_layers→mixing_layers(withso2_layerskept as an alias).