refactor: extract shared tabulate math, remove pt_expt→pt dependency#5339
Open
wanghan-iapcm wants to merge 11 commits intodeepmodeling:masterfrom
Open
refactor: extract shared tabulate math, remove pt_expt→pt dependency#5339wanghan-iapcm wants to merge 11 commits intodeepmodeling:masterfrom
wanghan-iapcm wants to merge 11 commits intodeepmodeling:masterfrom
Conversation
added 4 commits
March 25, 2026 15:06
Backend-agnostic DPTabulate with numpy implementations of: - Activation derivatives (grad, grad_grad) - Chain-rule propagation (unaggregated_dy_dx_s, etc.) - Embedding net forward pass (_make_data, _layer_0, _layer_1) - Network variable extraction (_get_network_variable) - Descriptor type detection from serialized data
Remove ~500 lines of duplicated math. PT DPTabulate now only overrides _get_descrpt_type (isinstance checks) and _convert_numpy_to_tensor (torch tensor conversion).
…dency DPTabulate now inherits from deepmd.utils.tabulate_math.DPTabulate instead of deepmd.pt.utils.tabulate.DPTabulate. Removed all imports of ActivationFn and DPTabulatePT from deepmd.pt. Descriptors now pass activation_fn_name (string) instead of ActivationFn (torch module).
Normalize lower/upper to per-type scalars after _get_env_mat_range so build() no longer needs is_pt branching. Move _convert_numpy_float_to_int into subclass _convert_numpy_to_tensor. Remove is_pt parameter from BaseTabulate and all callers.
![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: 56bdcf77a9
ℹ️ 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".
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRefactors tabulation math into a new numpy-only backend module and replaces PT-side math with a thin PyTorch wrapper that forwards activation names and tensor conversion; several PT call sites now pass raw activation names instead of PT Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
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)
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: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/utils/tabulate_math.py`:
- Around line 33-42: ACTIVATION_TO_FUNCTYPE currently omits aliases like
"linear" and "none" and fails to normalize prefixed names such as "silut*" and
"custom_silu*", causing lookups (where the code maps an activation name to
ACTIVATION_TO_FUNCTYPE) to raise RuntimeError; update ACTIVATION_TO_FUNCTYPE to
include "linear" and "none" entries and extend it with any derivative names you
need, and modify the lookup code to normalize inputs (e.g., lowercasing and
stripping/normalizing known prefixes like "silut", "custom_silu", "silu", etc.)
before dict lookup; also add matching derivative handling cases in the grad and
grad_grad implementations (ensure they accept the same normalized
aliases/prefixes) so the activation, grad, and grad_grad mappings are
consistent.
- Around line 438-440: The code currently assigns vv = zz.astype(self.data_type)
but zz is only defined for layer > 0, causing UnboundLocalError for single-layer
nets; change the assignment to use the current-layer output variable yy (i.e.,
set vv from yy.astype(self.data_type)) and ensure corresponding dd and d2 still
come from dy.astype(self.data_type) and dy2.astype(self.data_type) so both
one-layer and multi-layer cases work (look for vv, zz, yy, dy, dy2, and
data_type to locate the lines to update).
- Around line 246-259: The __init__ signature currently uses a mutable default
exclude_types: list[list[int]] = [], which causes a shared list across
instances; change the parameter to exclude_types: list[list[int]] | None = None
and inside DPTabulate.__init__ (before calling super().__init__) set
exclude_types = [] if exclude_types is None so each instance gets its own list,
then pass that local variable to super().__init__.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: fa001961-2e88-4010-890f-763ebb944a9c
📒 Files selected for processing (11)
deepmd/pt/utils/tabulate.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/se_e2_a.pydeepmd/pt_expt/descriptor/se_r.pydeepmd/pt_expt/descriptor/se_t.pydeepmd/pt_expt/descriptor/se_t_tebd.pydeepmd/pt_expt/utils/tabulate.pydeepmd/tf/utils/tabulate.pydeepmd/utils/tabulate.pydeepmd/utils/tabulate_math.py
💤 Files with no reviewable changes (1)
- deepmd/tf/utils/tabulate.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/utils/tabulate.py (1)
57-65:⚠️ Potential issue | 🟡 MinorFix shared defaults and normalize activation input type in
__init__.Lines 57–58 use shared mutable default objects (a code smell), and the type signature at line 58 is misleading since the parameter can reasonably accept strings at call sites. Line 65's direct access to
.activationassumes a specific type, which is fragile if the type contract isn't enforced.The proposed fix correctly:
- Replaces mutable defaults with
Noneand creates fresh objects- Widens the type annotation to accept both
strandActivationFn- Normalizes the input before passing to the parent class
Apply the suggested diff to improve type safety and prevent accidental state sharing across instances.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/utils/tabulate.py` around lines 57 - 65, The __init__ currently uses a shared mutable default for exclude_types and assumes activation_fn is an ActivationFn instance; change the signature to accept exclude_types: Optional[list[list[int]]] = None and activation_fn: Union[str, ActivationFn] = "tanh", then inside __init__ set exclude_types = [] if exclude_types is None and normalize activation_value = activation_fn.activation if isinstance(activation_fn, ActivationFn) else str(activation_fn) before calling super().__init__(..., exclude_types, activation_fn_name=activation_value) to avoid shared state and handle string inputs safely.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pt/utils/tabulate.py`:
- Around line 2-7: The constructor currently uses mutable and call-in-defaults
which fail linters: change the parameter defaults in the DPTabulate (or relevant
__init__) signature from exclude_types: list[list[int]] = [] to exclude_types:
Optional[list[list[int]]] = None and from activation_fn: ActivationFn =
ActivationFn("tanh") to activation_fn: Optional[ActivationFn] = None; then
inside the constructor, if exclude_types is None set exclude_types = [] (or a
fresh list) and if activation_fn is None set activation_fn =
ActivationFn("tanh"). Update any type hints/imports to include Optional and
ensure the instance attribute assignments use these local variables instead of
the original defaults.
---
Outside diff comments:
In `@deepmd/pt/utils/tabulate.py`:
- Around line 57-65: The __init__ currently uses a shared mutable default for
exclude_types and assumes activation_fn is an ActivationFn instance; change the
signature to accept exclude_types: Optional[list[list[int]]] = None and
activation_fn: Union[str, ActivationFn] = "tanh", then inside __init__ set
exclude_types = [] if exclude_types is None and normalize activation_value =
activation_fn.activation if isinstance(activation_fn, ActivationFn) else
str(activation_fn) before calling super().__init__(..., exclude_types,
activation_fn_name=activation_value) to avoid shared state and handle string
inputs safely.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 61f3ebed-df39-41bc-bd39-aeaaf52ceaad
📒 Files selected for processing (1)
deepmd/pt/utils/tabulate.py
… tensors The derivative functions now return numpy arrays from tabulate_math. Update test inputs/outputs accordingly. Remove unused torch/env imports.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5339 +/- ##
==========================================
- Coverage 82.28% 82.27% -0.01%
==========================================
Files 797 798 +1
Lines 82100 82110 +10
Branches 4003 4003
==========================================
+ Hits 67557 67558 +1
- Misses 13336 13345 +9
Partials 1207 1207 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
added 2 commits
March 25, 2026 21:57
- Fix UnboundLocalError: use yy instead of zz for single-layer nets - Add linear/none (functype=0) to ACTIVATION_TO_FUNCTYPE and grad/grad_grad (identity: f'=1, f''=0) - All 8 activation derivatives verified against numerical differentiation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
source/tests/pt/test_tabulate.py (2)
30-34: Unused dictionaryACTIVATION_NAMES_NUMPY_ONLY.This dictionary is defined but never referenced in the test code. Consider removing it or using it in
test_linear_activation.♻️ Option 1: Remove unused dict
-# functype=0 (linear/none) is not supported by TF custom ops, -# so we test it separately against numerical derivatives. -ACTIVATION_NAMES_NUMPY_ONLY = { - 0: "linear", -}♻️ Option 2: Use it in test_linear_activation
def test_linear_activation(self) -> None: - """Test functype=0 (linear/none) against numerical derivatives. + """Test numpy-only activations against numerical derivatives. TF custom ops don't support functype=0, so we validate against finite-difference derivatives instead. """ from deepmd.utils.tabulate_math import ( grad, grad_grad, ) - fn = get_activation_fn("linear") + for functype, activation_name in ACTIVATION_NAMES_NUMPY_ONLY.items(): + fn = get_activation_fn(activation_name)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/test_tabulate.py` around lines 30 - 34, The dictionary ACTIVATION_NAMES_NUMPY_ONLY is defined but never used; either remove it or make test_linear_activation consume it: update the test_linear_activation function to iterate or parameterize over ACTIVATION_NAMES_NUMPY_ONLY (using its keys/values) so the "linear" case is covered via that mapping, or simply delete the ACTIVATION_NAMES_NUMPY_ONLY constant if you prefer to keep the test hardcoded—ensure references to ACTIVATION_NAMES_NUMPY_ONLY are removed if deleting.
205-207: Unused variableh.The variable
h = 1e-7appears to be a step size for finite-difference derivatives, but the test uses analytical derivatives instead. Remove it or add the numerical verification it was intended for.♻️ Proposed fix - remove unused variable
fn = get_activation_fn("linear") y = fn(self.xbar) - h = 1e-7 # grad: f'(x) = 1 for identity🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/test_tabulate.py` around lines 205 - 207, Remove the unused numerical step variable h = 1e-7 from the test (or alternatively implement the intended finite-difference check); specifically, edit the test around get_activation_fn("linear") / fn and y = fn(self.xbar) to either delete the line defining h or replace it with a numeric-derivative verification that uses h to compute finite-difference gradients and compare them to the analytical derivatives returned by the activation function.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/utils/tabulate_math.py`:
- Around line 80-107: grad_grad currently returns -np.ones_like(xbar) for
unknown functype while grad raises a ValueError; update grad_grad so it mirrors
grad's behavior by raising a ValueError for unsupported functype values (include
the invalid functype in the exception message) instead of returning a sentinel
array; locate the grad_grad function and replace the final else branch to raise
the same style of ValueError used by grad.
- Around line 458-473: The function _get_descrpt_type builds a type_map that
lacks the "dpa2" => "Atten" entry, causing RuntimeError for DPA2 descriptors;
update the type_map inside _get_descrpt_type (symbol: type_map in
deepmd/utils/tabulate_math.py) to include "dpa2": "Atten" so DPA2 serializations
map to "Atten", and apply the identical change to the analogous mapping in
deepmd/pt_expt/utils/tabulate.py (around the mapping at line 65) to keep both
modules consistent.
---
Nitpick comments:
In `@source/tests/pt/test_tabulate.py`:
- Around line 30-34: The dictionary ACTIVATION_NAMES_NUMPY_ONLY is defined but
never used; either remove it or make test_linear_activation consume it: update
the test_linear_activation function to iterate or parameterize over
ACTIVATION_NAMES_NUMPY_ONLY (using its keys/values) so the "linear" case is
covered via that mapping, or simply delete the ACTIVATION_NAMES_NUMPY_ONLY
constant if you prefer to keep the test hardcoded—ensure references to
ACTIVATION_NAMES_NUMPY_ONLY are removed if deleting.
- Around line 205-207: Remove the unused numerical step variable h = 1e-7 from
the test (or alternatively implement the intended finite-difference check);
specifically, edit the test around get_activation_fn("linear") / fn and y =
fn(self.xbar) to either delete the line defining h or replace it with a
numeric-derivative verification that uses h to compute finite-difference
gradients and compare them to the analytical derivatives returned by the
activation function.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 07fbed44-6815-422f-8d2e-6b756257a489
📒 Files selected for processing (2)
deepmd/utils/tabulate_math.pysource/tests/pt/test_tabulate.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
deepmd/utils/tabulate_math.py (2)
283-286:⚠️ Potential issue | 🟠 MajorMake the DPA2 branch reachable.
Lines 283-286 explicitly handle DPA2 by switching to
serialized["repinit_variable"], but_get_descrpt_type()never maps"dpa2"to"Atten", so serialization-driven callers fail before that branch can run. Add"dpa2": "Atten"to the type map.🐛 Proposed fix
type_map = { "se_e2_a": "A", "se_r": "R", "se_e3": "T", "se_e3_tebd": "T_TEBD", "dpa1": "Atten", + "dpa2": "Atten", "se_atten_v2": "Atten", }Also applies to: 463-473
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/utils/tabulate_math.py` around lines 283 - 286, The DPA2 branch isn't reachable because _get_descrpt_type() does not map "dpa2" to "Atten", so the later code that expects descrpt_type == "Atten" and uses serialized["repinit_variable"] never runs; fix by adding the mapping "dpa2": "Atten" to the type map inside _get_descrpt_type() (and mirror this change in the other similar mapping block referenced around the other occurrence), ensuring callers that serialize can hit the serialized["repinit_variable"] branch.
33-44:⚠️ Potential issue | 🟠 MajorNormalize activation aliases before the functype lookup.
The shared backend now accepts the forward activation via
get_activation_fn(...), but the derivative side still only accepts exact keys fromACTIVATION_TO_FUNCTYPE. That leaves valid project aliases/prefixed SiLU names failing here withUnknown activation function ...even though the forward activation can still be constructed. Normalize once up front and use that normalized name for bothget_activation_fn()and the functype map.♻️ Proposed fix
+def _normalize_activation_name(name: str) -> str: + normalized = name.lower() + if normalized.startswith(("silut", "custom_silu")): + return "silu" + if normalized in {"none", "linear"}: + return "linear" + return normalized + ... - self._activation_fn = get_activation_fn(activation_fn_name) - activation_fn_name = activation_fn_name.lower() + activation_fn_name = _normalize_activation_name(activation_fn_name) + self._activation_fn = get_activation_fn(activation_fn_name) if activation_fn_name not in ACTIVATION_TO_FUNCTYPE: raise RuntimeError(f"Unknown activation function: {activation_fn_name}")Also applies to: 267-271
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/utils/tabulate_math.py` around lines 33 - 44, The ACTIVATION_TO_FUNCTYPE lookup fails for valid aliases because activation names are not normalized first; normalize the activation string once (using the same normalization logic as get_activation_fn or by calling a canonicalizer used by get_activation_fn) and then use that normalized name for both calling get_activation_fn(...) and for the ACTIVATION_TO_FUNCTYPE[...] lookup; apply the same change at the other occurrence that also performs a functype lookup so both sites use the same canonical activation key.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@deepmd/utils/tabulate_math.py`:
- Around line 283-286: The DPA2 branch isn't reachable because
_get_descrpt_type() does not map "dpa2" to "Atten", so the later code that
expects descrpt_type == "Atten" and uses serialized["repinit_variable"] never
runs; fix by adding the mapping "dpa2": "Atten" to the type map inside
_get_descrpt_type() (and mirror this change in the other similar mapping block
referenced around the other occurrence), ensuring callers that serialize can hit
the serialized["repinit_variable"] branch.
- Around line 33-44: The ACTIVATION_TO_FUNCTYPE lookup fails for valid aliases
because activation names are not normalized first; normalize the activation
string once (using the same normalization logic as get_activation_fn or by
calling a canonicalizer used by get_activation_fn) and then use that normalized
name for both calling get_activation_fn(...) and for the
ACTIVATION_TO_FUNCTYPE[...] lookup; apply the same change at the other
occurrence that also performs a functype lookup so both sites use the same
canonical activation key.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: c5ea7a07-c5f4-489a-b10d-cbf14a56293f
📒 Files selected for processing (4)
deepmd/pt/utils/tabulate.pydeepmd/pt_expt/utils/tabulate.pydeepmd/utils/tabulate_math.pysource/tests/pt/test_tabulate.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/pt/test_tabulate.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
1 participant
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.
Summary
deepmd/pt/utils/tabulate.pyinto a shared numpy moduledeepmd/utils/tabulate_math.pyDPTabulateclassdeepmd.ptfor tabulation or compressionis_ptflag fromBaseTabulate— normalize lower/upper bounds to per-type scalars, move_convert_numpy_float_to_intinto subclass_convert_numpy_to_tensorChanges
deepmd/utils/tabulate_math.pydeepmd/pt/utils/tabulate.pydeepmd/pt_expt/utils/tabulate.pydeepmd/pt_expt/descriptor/*.pydeepmd/utils/tabulate.pydeepmd/tf/utils/tabulate.pyTest plan
deepmd.ptimports in pt_expt tabulate or descriptorsis_ptreferences remainSummary by CodeRabbit
Refactor
New Features
Tests