Skip to content

Commit 51ffe61

Browse files
wanghan-iapcmHan Wang
andauthored
refactor: extract shared tabulate math, remove pt_expt→pt dependency (#5339)
## Summary - Extract tabulation math (activation derivatives, chain-rule propagation, embedding-net forward pass) from `deepmd/pt/utils/tabulate.py` into a shared numpy module `deepmd/utils/tabulate_math.py` - Both pt and pt_expt backends now inherit from the shared `DPTabulate` class - pt_expt no longer imports from `deepmd.pt` for tabulation or compression - Remove `is_pt` flag from `BaseTabulate` — normalize lower/upper bounds to per-type scalars, move `_convert_numpy_float_to_int` into subclass `_convert_numpy_to_tensor` ## Changes | File | Lines | What | |------|-------|------| | `deepmd/utils/tabulate_math.py` | +582 | New: numpy DPTabulate with all math | | `deepmd/pt/utils/tabulate.py` | -560 | Thin wrapper: _get_descrpt_type + _convert_numpy_to_tensor | | `deepmd/pt_expt/utils/tabulate.py` | rewrite | Inherit from tabulate_math, no pt imports | | `deepmd/pt_expt/descriptor/*.py` | -6 files | Remove ActivationFn import, pass string | | `deepmd/utils/tabulate.py` | -12 | Remove is_pt flag and branching | | `deepmd/tf/utils/tabulate.py` | -1 | Remove is_pt=False arg | ## Test plan - [x] PT compression tests pass (16 tests: se_a, se_atten, se_r, se_t, dpa2) - [x] PT_EXPT model compression tests pass (6 tests) - [x] PT_EXPT descriptor compression tests pass (16 tests) - [x] No `deepmd.pt` imports in pt_expt tabulate or descriptors - [x] No `is_pt` references remain <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Streamlined tabulation: removed backend-specific PT logic, normalized activation/exclude handling, and standardized numpy→tensor conversion into a thin runtime wrapper. * **New Features** * Added a backend‑neutral math tabulation component that provides activation derivatives and chain‑rule propagation for embedding computations. * **Tests** * Converted tests to NumPy inputs/outputs for backend independence and added an explicit linear‑activation derivative test. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 2a82988 commit 51ffe61

13 files changed

Lines changed: 784 additions & 676 deletions

File tree

deepmd/dpmodel/utils/network.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,18 @@ def sigmoid_t(x): # noqa: ANN001, ANN201
4444
return xp_sigmoid(x)
4545

4646

47+
def softplus_t(x): # noqa: ANN001, ANN201
48+
"""Numerically stable softplus."""
49+
xp = array_api_compat.array_namespace(x)
50+
positive = x > 0
51+
exp_neg_abs = xp.exp(xp.where(positive, -x, x))
52+
return xp.where(
53+
positive,
54+
x + xp.log(1 + exp_neg_abs),
55+
xp.log(1 + exp_neg_abs),
56+
)
57+
58+
4759
class Identity(NativeOP):
4860
def __init__(self) -> None:
4961
super().__init__()
@@ -330,9 +342,7 @@ def fn(x): # noqa: ANN001, ANN202
330342
elif activation_function == "softplus":
331343

332344
def fn(x): # noqa: ANN001, ANN202
333-
xp = array_api_compat.array_namespace(x)
334-
# generated by GitHub Copilot
335-
return xp.log(1 + xp.exp(x))
345+
return softplus_t(x)
336346

337347
return fn
338348
elif activation_function == "sigmoid":

0 commit comments

Comments
 (0)