docs(backend): document backend support and workflows - #5929
Conversation
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
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:
📝 WalkthroughWalkthroughDocumentation now covers seven backends, including TensorFlow 2 and PyTorch-Exportable, across backend selection, training, freezing, compression, model formats, inference detection, and Hessian handling. ChangesBackend documentation
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@doc/train/training.md`:
- Line 23: Remove the shell prompt marker from each command-only example, or add
corresponding output, to eliminate MD014 warnings: update the TensorFlow 2,
PyTorch, and PyTorch-Exportable examples at doc/train/training.md lines 23, 30,
and 37, and the TensorFlow 2 and PyTorch-Exportable examples at
doc/freeze/freeze.md lines 21 and 52.
🪄 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 Plus
Run ID: 16cb4b9c-60c4-416c-bbc6-ec382d110acd
📒 Files selected for processing (4)
doc/backend.mddoc/freeze/compress.mddoc/freeze/freeze.mddoc/train/training.md
There was a problem hiding this comment.
Pull request overview
This PR adds user-facing documentation for selecting and using the TensorFlow 2 (eager) and PyTorch-Exportable (pt-expt) backends, covering backend selection plus train, freeze, and compress CLI usage across the relevant docs pages.
Changes:
- Add TF2 and PyTorch-Exportable tabs to the training instructions.
- Add TF2 and PyTorch-Exportable tabs to the freeze and compress docs (including
--lower-kindmention for pt-expt). - Extend
doc/backend.mdto include the two backends, their model/checkpoint formats, and inference suffix detection notes.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| doc/backend.md | Adds TF2 and PyTorch-Exportable backend documentation (selection, file formats, detection). |
| doc/train/training.md | Adds TF2 and PyTorch-Exportable training command tabs. |
| doc/freeze/freeze.md | Adds TF2 and PyTorch-Exportable freeze command tabs and notes. |
| doc/freeze/compress.md | Updates supported-backends note and adds TF2 / PyTorch-Exportable compress tabs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Document TF2, PyTorch-Exportable, and JAX availability across model and descriptor guides, and add DP/JAX compression routes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
njzjz
changed the title
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)
doc/backend.md (1)
114-117: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDocument the DPA4
.pt2ABI, not just the filename classification.
doc/model/dpa4.mddocumentsdp --pt freezefor DPA4 anddoc/backend.mdsays.pt2belongs to PyTorch-Exportable. The code docs show thesept.pt2artifacts are AOTInductor, but both ABIs can be classified as.pt2, while the PT and pt_expt code paths/ABI are distinct. Make the inference guidance explicit about “which.pt2” or rename/correct one contract to avoid loading the wrong graph type.🤖 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 `@doc/backend.md` around lines 114 - 117, Update the `.pt2` inference guidance in the backend documentation to distinguish DPA4 `dp --pt freeze` artifacts from PyTorch-Exportable/AOTInductor artifacts, explicitly identifying which ABI each `.pt2` file uses. Align the wording with the separate PT and pt_expt code paths so users do not load a `.pt2` graph through the wrong backend contract.
🤖 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 `@doc/model/train-fitting-dos.md`:
- Around line 4-11: Update the DOS training documentation to cover TensorFlow 2
and PyTorch-Exportable routes alongside the existing dp --tf train and dp --pt
train examples, adding dp --tf2 train and dp --pt-expt train tabs. If either
backend cannot train DOS models, instead revise the supported-backends
description to explicitly mark it as conversion/inference-only, consistent with
JAX.
---
Outside diff comments:
In `@doc/backend.md`:
- Around line 114-117: Update the `.pt2` inference guidance in the backend
documentation to distinguish DPA4 `dp --pt freeze` artifacts from
PyTorch-Exportable/AOTInductor artifacts, explicitly identifying which ABI each
`.pt2` file uses. Align the wording with the separate PT and pt_expt code paths
so users do not load a `.pt2` graph through the wrong backend contract.
🪄 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 Plus
Run ID: 2b1880b6-d0b4-460d-b54d-d31775212ee3
📒 Files selected for processing (24)
doc/backend.mddoc/freeze/compress.mddoc/freeze/freeze.mddoc/model/dpa2.mddoc/model/dpa3.mddoc/model/dpa4.mddoc/model/dprc.mddoc/model/linear.mddoc/model/overall.mddoc/model/pairtab.mddoc/model/train-energy-hessian.mddoc/model/train-energy-spin.mddoc/model/train-energy.mddoc/model/train-fitting-dos.mddoc/model/train-fitting-property.mddoc/model/train-fitting-tensor.mddoc/model/train-hybrid.mddoc/model/train-se-atten.mddoc/model/train-se-e2-a-tebd.mddoc/model/train-se-e2-a.mddoc/model/train-se-e2-r.mddoc/model/train-se-e3-tebd.mddoc/model/train-se-e3.mddoc/train/training.md
🚧 Files skipped from review as they are similar to previous changes (1)
- doc/freeze/compress.md
Document the distinct .pt2 lower-input ABIs and add the missing DOS training commands for TensorFlow 2 and PyTorch-Exportable. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Document the dense nlist exception used by PT-frozen deepspin DPA4 models. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5929 +/- ##
==========================================
- Coverage 79.59% 79.34% -0.26%
==========================================
Files 1081 1081
Lines 126244 126244
Branches 4592 4598 +6
==========================================
- Hits 100490 100173 -317
- Misses 24101 24419 +318
+ Partials 1653 1652 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
This is a large and mostly accurate sweep -- I checked the great majority of the new support claims against backend registration at head and they hold up, including the ones that are easy to get wrong. Specifically confirmed correct: every dp --tf2 / dp --pt-expt / dp --jax flag against the aliases in deepmd/backend/*.py; the dp --pt_expt to dp --pt-expt and .xlo to .hlo corrections (both were real doc bugs); all the frozen-model suffixes against each backend's suffixes; dp --jax freeze --hessian; the JAX energy-loss-only and no---init-frz-model caveats; the TF2 and Paddle additions on dpa2/dpa3/se_e3_tebd; and the DPA4 note, where JAX really does have deepmd/jax/descriptor/dpa4.py but no jax/model/dpa4_model.py, so "descriptor-only" is exactly right.
Two claims do not hold. Both are inline.
One smaller thing I did not raise inline. The Hessian rewording in overall.md and train-energy-hessian.md is not wrong -- it names dp --pt freeze and the TorchScript route specifically -- but hessian_mode is honored only at deepmd/pt/model/model/__init__.py#L371-L374, and neither pt_expt/model/get_model.py nor dpmodel/model/model.py mention it, so PyTorch-Exportable cannot construct a Hessian model at all. Since this PR establishes that {{ pytorch_icon }} covers both, the page is now silent on pt_expt where it used to be blanket. Worth a clause if you touch these lines again; not worth a round trip on its own.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The current DPA4 native-spin multi-rank guidance contradicts both the implementation and itself. The graph-route/native-spin section of doc/model/dpa4.md says that native-spin .pt2 archives carry no with-comm artifact, have has_comm_artifact=false, and are single-rank only. At this head, _needs_with_comm_artifact explicitly admits a graph-lowered NativeSpinModelKind whose descriptor communicates across ranks; test_native_spin_needs_with_comm_on_the_graph_lower_only asserts that behavior; and DeepSpinPTExpt loads and dispatches forward_lower_with_comm.pt2 for multi-rank graph native spin. The same page also correctly states near the earlier multi-rank overview and the Limitations section that pt_expt native spin supports multi-rank. Please remove or update the obsolete single-rank paragraphs so this page documents one contract that matches the current export and C++ runtime.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The documented PyTorch-Exportable compression workflow does not actually put a compressed model into the executable artifact for the ordinary .pte example. Except for the special fused DPA1 graph .pt2 branch, enable_compression passes the original uncompressed model dictionary to deserialize_to_file and stores the compressed dictionary only as model_json_override. DeepEval reconstructs that model.json object for metadata and introspection, but energy inference runs exported_module or _pt2_runner, so the tabulated model is not executed. I reproduced the documented se_e2_a-style .pte workflow locally: the input and output exported graphs were identical (711 nodes each), the output executable contained no tabulate_fusion operator, while only model.json contained the compression state. Therefore model-compress.pte provides no compression speed or memory benefit even though this page and backend.md present pt_expt compression as generally supported. Please either export the genuinely compressed executable or document the narrow supported case instead of advertising the ordinary .pte workflow as compressed.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The newly documented JAX .jax frozen-model route is not usable by the normal DeePMD inference/test interface. Please correct the format contract or implement the missing loader support.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The newly documented TensorFlow 2 compression workflow fails for the normalized default se_e2_a configuration. The implementation or the documented support contract needs to be corrected.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One remaining backend-classification error is noted inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional backend-workflow omission is noted inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional PyTorch-Exportable training-workflow omission is noted inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional graph-freeze behavior mismatch is noted inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional backend-workflow documentation gap is noted inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional spin-scheme semantic mismatch is noted inline.
# Conflicts: # doc/model/train-se-atten.md
Reply and close the current review round on deepmodeling#5929: - backend.md: describe the DPA4/SeZM `dp --pt freeze` export as the separate AOTInductor route instead of a PyTorch-TorchScript route, so the overview agrees with dpa4.md. - change-bias.md: document the PyTorch-Exportable `dp --pt-expt change-bias` route for .pt/.pte/.pt2 inputs, with the multi-task limitation noted. - parallel-training.md: add the PyTorch-Exportable DDP subsection with a `torchrun ... dp --pt-expt train` example, separate from the PyTorch-TorchScript ZeRO/FSDP2 section. - dpa2.md: drop the unavailable dense (`nlist`, default) fallback advice and state the forced graph lower for graph-eligible DPA-2. - finetuning.md: document the TensorFlow 2 and PyTorch-Exportable `--finetune` workflows and their accepted formats/limitations. - train-energy-spin.md: scope the virtual-atom (sel-extension) guidance to the deepspin scheme, add the DPA4 native exception, and add DPA4 to the scheme-aware supported-descriptor list. Coding-Agent: opencode opencode-Version: 1.18.11 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
|
Addressed the current round of review feedback in 708891e and merged latest master (resolving the merge conflict):
All review threads replied to inline and resolved. Validation: mdformat passed on all changed files, Coding agent: opencode |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.
Suppressed comments (1)
doc/freeze/compress.md:122
- The DP compression note omits the
.ymlDPModel format.deepmd.dpmodel.utils.serialization.load_dp_model()explicitly supports both.yamland.yml, so the docs should list both to avoid implying.ymlinputs are unsupported.
DP compression accepts native `.dp` and `.yaml` models.
- compress.md: list `.yml` alongside `.dp`/`.yaml` for DP compression, matching load_dp_model's supported extensions. - train-se-e2-r.md: include PyTorch-Exportable in the backends for which `type_one_side=false` is unsupported (pt_expt `se_r` inherits the dpmodel implementation that raises NotImplementedError). - howtoset_num_nodes.md: note that torchrun-based multiprocessing training applies to both PyTorch-TorchScript and PyTorch-Exportable. Coding-Agent: opencode opencode-Version: 1.18.11 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.
Suppressed comments (3)
doc/train/finetuning.md:5
- The supported-backends note omits PyTorch-Exportable even though this page now documents a PyTorch-Exportable fine-tuning workflow below. This is misleading for readers scanning the page header; list both PyTorch variants here for consistency with the rest of the backend docs.
> **Supported backends**: TensorFlow {{ tensorflow_icon }}, TensorFlow 2
> {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}
doc/model/train-energy-hessian.md:118
- This warning states that PyTorch-Exportable cannot construct a Hessian model for freezing or inference, but the codebase includes explicit pt_expt Hessian-model support (e.g.
deepmd/pt_expt/model/ener_model.pyhasenable_hessian()and there are pt_expt Hessian tests). The doc should not make an absolute “cannot construct” claim here; at minimum, remove the pt-expt sentence or replace it with a precise, verified limitation about export/freeze behavior.
> The PyTorch-TorchScript freeze route does not preserve Hessian output. A
> PyTorch-TorchScript model frozen with `dp --pt freeze` is treated as a
> standard energy model. PyTorch-Exportable cannot construct a Hessian model
> for freezing or inference. The JAX backend can preserve Hessian output in a
> frozen model with `dp --jax freeze --hessian`.
doc/model/overall.md:64
- This warning claims PyTorch-Exportable cannot construct a Hessian model for freezing or inference, but pt_expt has Hessian-model support in code/tests (e.g.
EnergyModel.enable_hessian()andsource/tests/pt_expt/model/test_ener_hessian_model.py). The documentation should not state that pt-expt cannot construct Hessian models; remove or rephrase this sentence to match the implemented behavior.
> The PyTorch-TorchScript freeze route cannot output Hessians, and
> PyTorch-Exportable cannot construct a Hessian model for freezing or
> inference. The JAX backend can retain Hessian output in a frozen model with
> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md).
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The native-spin section still documents a lower-kind choice that the current pt_expt freeze interface does not provide. It says that a plain-energy DPA4/SeZM descriptor can freeze to either the dense or graph lower and that --lower-kind auto is the default. However, the CLI accepts only nlist and graph, defaults to nlist, and freeze() overrides that default/request to graph whenever model_uses_graph_lower(m) is true. Since an ordinary DPA4 descriptor reports uses_graph_lower() == True, it is graph-frozen as well; auto is not a valid value. Please make this paragraph match the current forced-graph policy and remove the nonexistent auto option.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The earlier correction that distinguishes JAX .jax serialization from DeepEval-compatible frozen formats is still incomplete in the Hessian documentation. train-energy-hessian.md says backend checkpoints can be tested directly and gives dp --jax test -m model.ckpt.jax, but test is a common entrypoint and the JAX DeepEval constructor accepts only .hlo and .savedmodel; a .jax path raises ValueError(Unsupported file extension). Please remove the direct .jax test workflow and use the already documented freeze --hessian to .hlo/.savedmodel route instead, unless .jax DeepEval loading is implemented.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The .pt2 lower-input ABI overview is incomplete: it enumerates edge_vec, nlist, and graph and then describes these as all variants, but the documented pt_expt DPA1 compression workflow can produce a fourth production ABI, lower_input_kind=dpa1_canonical, whenever canonical_model_eligible(model) is true. DeepPotPTExpt also has a distinct lower_input_is_canonical_ dispatch branch for it. Please document dpa1_canonical as the compact compressed-DPA1 ABI in both the ABI overview and the later extension-based inference-routing summary instead of subsuming it under NeighborGraph.
Address the remaining CHANGES_REQUESTED round on deepmodeling#5929: - dpa4.md: remove the nonexistent --lower-kind auto option; state that the CLI accepts only nlist (default) and graph, that a graph-capable DPA4/SeZM model is always frozen to the graph lower (nlist is overridden with a warning), and that a native-spin descriptor has no dense lower at all. - train-energy-hessian.md: drop the direct dp --jax test of .jax checkpoints, which the JAX DeepEval route does not accept; JAX models must be frozen to .hlo/.savedmodel with dp --jax freeze --hessian before testing, as already documented above. - backend.md: document the compact dpa1_canonical ABI produced by dp --pt-expt compress for eligible compressed DPA-1 in the .pt2 lower-input ABI overview and the inference-routing summary. Coding-Agent: opencode opencode-Version: 1.18.13 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
All three findings from the latest CHANGES_REQUESTED round are addressed in commit e095aa7e6:
-
dpa4.md (
--lower-kind auto) — Removed the nonexistentautooption. The text now states--lower-kindaccepts onlynlist(the default) andgraph, that a graph-capable DPA4/SeZM descriptor (plain-energy or native-spin) is always frozen to the graph lower because the freeze entrypoint overrides annlistrequest with a warning, and that a native-spin descriptor has no dense (nlist) lower at all, so--lower-kindcannot produce a dense artifact for it. -
train-energy-hessian.md (
dp --jax test -m model.ckpt.jax) — Removed the direct.jaxtest workflow (JAXdp testaccepts only frozen.hlo/.savedmodel). The page now keeps the PyTorch-TorchScript direct-checkpoint test and states that a JAX Hessian model must be frozen withdp --jax freeze --hessianfirst, pointing to the freeze route documented above. -
backend.md (
.pt2lower-input ABI overview) — Documented the fourth production ABI:dp --pt-expt compresscan export an eligible compressed DPA-1 model through the compact canonicaldpa1_canonicalABI (lower_input_kind: dpa1_canonical). It is now listed both in the.pt2ABI overview and in the extension-based inference-routing summary alongsideedge_vec, densenlist, and NeighborGraph.
Coding agent: opencode
opencode version: 1.18.13
Model: ustc/deepseek-v4-flash
Reasoning effort: max
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.
Suppressed comments (3)
doc/train/finetuning.md:5
- The supported-backends note omits PyTorch-Exportable even though this page now documents a PyTorch-Exportable fine-tuning workflow below. This is misleading because it suggests pt-expt fine-tuning is out of scope for this page.
> [!NOTE]
> **Supported backends**: TensorFlow {{ tensorflow_icon }}, TensorFlow 2
> {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}
doc/model/train-energy-hessian.md:165
- This claims
.jaxcheckpoints cannot be tested directly, but the JAX backend is registered for.jaxmodel paths and the DeepEval implementation can load.jaxvia Orbax (serialize_from_fileindeepmd/jax/utils/serialization.py). As written, it incorrectly forbids a workflow that should work (dp --jax test -m model.ckpt.jax ...).
The JAX training checkpoint (`.jax`) cannot be tested directly: the JAX
`dp test` route accepts only frozen `.hlo` or `.savedmodel` artifacts, so a
JAX Hessian model must be frozen with `dp --jax freeze --hessian` (as shown
above) before testing.
doc/backend.md:107
- This says
dp test/DeepPotdoes not load.jaxserializations, but.jaxis a registered JAX model suffix (deepmd/backend/jax.py) andDeepEvalBackenddispatches by suffix, sodp test -m *.jax/DeepPot('*.jax')should work. If the intent is to discourage.jaxfor deployment, it’s more accurate to scope the limitation to the C++ interface.
JAX supports training with `dp --jax train`; training checkpoints use the
`.jax` extension. Freezing can write a DeepEval-compatible `.hlo` or
`.savedmodel` model, or a lossless `.jax` serialization for checkpoint
round-tripping and JAX-MD. The normal `dp test`/`DeepPot` route does not load
`.jax` serializations.
# Conflicts: # doc/model/dpa4.md
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 33 out of 33 changed files in this pull request and generated no new comments.
Suppressed comments (3)
doc/backend.md:52
- This sentence asserts a future deprecation plan ("DeePMD-kit will deprecate this backend") without any pointer to a deprecation notice, release note, or tracking issue. In user docs this reads like a guarantee/timeline and can become stale quickly. Consider rephrasing to describe the current technical difference and, if desired, recommend pt-expt without promising deprecation.
[PyTorch](https://pytorch.org/) 2.1 or above is required. Select this backend
with `dp --pt`. It uses TorchScript for most frozen models; DPA4/SeZM uses a
separate AOTInductor export path. Because PyTorch has deprecated TorchScript,
DeePMD-kit will deprecate this backend and replace it with PyTorch-Exportable.
doc/model/train-energy-hessian.md:118
- The warning says “PyTorch-Exportable cannot construct a Hessian model for freezing or inference”, but the pt-expt implementation does support Hessian computation in eager mode (it is just not exportable / not preserved in frozen artifacts). To avoid contradicting the implementation, scope this statement to exported
.pte/.pt2models (torch.export/AOTI) and/or todp teston frozen artifacts.
> The PyTorch-TorchScript freeze route does not preserve Hessian output. A
> PyTorch-TorchScript model frozen with `dp --pt freeze` is treated as a
> standard energy model. PyTorch-Exportable cannot construct a Hessian model
> for freezing or inference. The JAX backend can preserve Hessian output in a
> frozen model with `dp --jax freeze --hessian`.
doc/model/overall.md:64
- Same issue as in
train-energy-hessian.md: the current wording is too broad (“PyTorch-Exportable cannot construct a Hessian model for freezing or inference”). The pt-expt backend can compute Hessians in eager mode, but Hessians are not preserved in exported.pte/.pt2artifacts. Consider scoping the warning accordingly so it matches the implementation.
> The PyTorch-TorchScript freeze route cannot output Hessians, and
> PyTorch-Exportable cannot construct a Hessian model for freezing or
> inference. The JAX backend can retain Hessian output in a frozen model with
> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md).
Summary
Closes #5919.
Closes #5920.
Closes #5923.
Supersedes #5930.
Validation
ruff format .ruff check .uvx pre-commit run mdformat --files ...for all changed Markdown filesdp --dp compress --helpdp --jax compress --helpgit diff --checkCoding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit