docs(backend): document backend support and workflows

doc/backend.md

@@ -2,10 +2,15 @@
 
 ## Supported backends
 
-DeePMD-kit supports multiple backends: TensorFlow and PyTorch.
+DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, PyTorch,
+PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP reference backend.
 To use DeePMD-kit, you must install at least one backend.
 Each backend does not support all features.
-In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, and Paddle {{ paddle_icon }} icons are used to mark whether a backend supports a feature.
+In the documentation, TensorFlow and TensorFlow 2 share
+{{ tensorflow_icon }}, while PyTorch and PyTorch-Exportable share
+{{ pytorch_icon }}. JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP
+{{ dpmodel_icon }} use separate icons. Support notes spell out the exact backend
+variant when the two implementations in a framework family differ.
 
 ### TensorFlow {{ tensorflow_icon }}
 
@@ -15,6 +20,20 @@ In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon
 [TensorFlow](https://tensorflow.org) 2.8 is the first version to support Python 3.10.
 DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`tf.compat.v1`) in the graph mode.
 
+### TensorFlow 2 {{ tensorflow_icon }}
+
+- Model filename extension: `.savedmodeltf`
+- Checkpoint directory extension: `.tf2`
+
+The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with
+`dp --tf2` (alias `dp --tensorflow2`). It supports training, including
+multi-task training and fine-tuning, freezing, compression, and testing.
+Training stores checkpoints in a directory named after the `save_ckpt` prefix
+with `.tf2` appended, such as `model.ckpt.tf2`.
+
+Setting [`DP_JIT`](env.md#envvar-DP_JIT) enables optional `tf.function` JIT
+compilation; depending on the workload, this may improve or reduce performance.
+
 ### PyTorch {{ pytorch_icon }}
 
 - Model filename extension: `.pth`
@@ -23,17 +42,39 @@ DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`t
 [PyTorch](https://pytorch.org/) 2.1 or above is required.
 While `.pth` and `.pt` are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint.
 
+### PyTorch-Exportable {{ pytorch_icon }}
+
+- Model filename extensions: `.pte`, `.pt2`
+- Checkpoint filename extension: `.pt`
+
+Select this backend with `dp --pt-expt` (alias
+`dp --pytorch-exportable`). It uses PyTorch with the backend-independent model
+implementation and supports training, including multi-task training and
+fine-tuning, freezing, compression, change-bias, and testing. Training can read
+LMDB datasets, and Python inference can use the optional vesin neighbor-list
+implementation.
+
+Freezing exports a `torch.export` model. The dense neighbor-list lower form
+normally uses `.pte`, while the graph lower form uses an AOTInductor `.pt2`
+package. Use `--lower-kind graph` to request graph-native export for an eligible
+model; graph-capable DPA models may select that form automatically. The `.pt`
+checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which
+allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint.
+
 ### JAX {{ jax_icon }}
 
-- Model filename extension: `.xlo`, `.savedmodel`
+- Model filename extensions: `.hlo`, `.jax`, `.savedmodel`
 - Checkpoint filename extension: `.jax`
 
 [JAX](https://jax.readthedocs.io/) 0.4.33 or above is required.
-Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
+Both `.hlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
 `.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow.
 Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface.
 The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs.
 
+JAX supports training with `dp --jax train`; training checkpoints use the
+`.jax` extension and can be frozen as `.hlo`, `.jax`, or `.savedmodel` models.
+
 ### Paddle {{ paddle_icon }}
 
 - Model filename extensions: `.json` and `.pdiparams`
@@ -64,12 +105,16 @@ NumPy 1.21 or above is required.
 
 ### Training
 
-When training and freezing a model, you can use `dp --tf`, `dp --pt` or `dp --pd` in the command line to switch the backend.
+When training and freezing a model, use `dp --tf`, `dp --tf2`, `dp --pt`,
+`dp --pt-expt`, `dp --jax`, or `dp --pd` in the command line to switch the
+backend.
 
 ### Inference
 
 When doing inference, DeePMD-kit detects the backend from the model filename.
 For example, when the model filename ends with `.pb` (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend.
+The same detection covers TensorFlow 2 `.savedmodeltf` models and
+PyTorch-Exportable `.pte` and `.pt2` models.
 
 ## Convert model files between backends
 

doc/freeze/compress.md

@@ -1,7 +1,9 @@
-# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }}
+# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}
+> **Backends covered below**: TensorFlow and TensorFlow 2
+> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX
+> {{ jax_icon }}, and DP {{ dpmodel_icon }}.
 
 ## Theory
 
@@ -70,19 +72,66 @@ dp compress -i graph.pb -o graph-compress.pb
 ```
 :::
 
+:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}
+
+```bash
+dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf
+```
+
+TensorFlow 2 compression reads a `.tf2` training checkpoint directory or a
+checkpoint prefix and writes a compressed `.savedmodeltf` model.
+:::
+
 :::{tab-item} PyTorch {{ pytorch_icon }}
 
 ```bash
 dp --pt compress -i model.pth -o model-compress.pth
 ```
 :::
 
+:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}
+
+```bash
+dp --pt-expt compress -i model.pte -o model-compress.pte
+```
+
+Use matching `.pte` or `.pt2` suffixes to preserve the exported model form.
+:::
+
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax compress -i frozen_model.hlo -o compressed_model.hlo
+```
+
+JAX compression accepts `.jax` and `.hlo` models. Use the same suffix for the
+input and output to preserve the serialized model format.
+:::
+
+:::{tab-item} DP {{ dpmodel_icon }}
+
+```bash
+dp --dp compress -i model.dp -o model-compress.dp
+```
+
+DP compression accepts native `.dp` and `.yaml` models.
+:::
+
 ::::
 
-where `-i` gives the original frozen model, `-o` gives the compressed model. Several other command line options can be passed to `dp compress`, which can be checked with
+where `-i` gives the original frozen model, `-o` gives the compressed model.
+The DP, JAX, and TensorFlow 2 entrypoints share the native-model compression
+helpers for resolving the minimum neighbor distance and tabulating the
+descriptor's embedding networks. PyTorch-Exportable implements the same table
+strides and minimum-neighbor-distance fallback in its export-specific
+entrypoint. If the model does not contain a minimum neighbor distance, pass the
+training script with `-t` or `--training-script` so it can be computed from the
+training data.
+
+Several other command line options can be passed to `dp compress`, which can be checked with
 
 ```bash
-$ dp compress --help
+dp compress --help
 ```
 
 An explanation will be provided

doc/freeze/freeze.md

@@ -8,17 +8,30 @@ To freeze a model, typically one does
 :::{tab-item} TensorFlow {{ tensorflow_icon }}
 
 ```bash
-$ dp freeze -o model.pb
+dp freeze -o model.pb
 ```
 
 in the folder where the model is trained. The output model is called `model.pb`.
 The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc).
 :::
 
+:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}
+
+```bash
+dp --tf2 freeze -c model.ckpt -o model.savedmodeltf
+```
+
+When `-c` names a checkpoint prefix, the backend also checks the corresponding
+path with `.tf2` appended, so the example reads `model.ckpt.tf2` and writes the
+TensorFlow SavedModel to `model.savedmodeltf`. If `-c` is omitted, it defaults
+to the current directory. For a multi-task checkpoint, select a branch with
+`--head CHOSEN_BRANCH`.
+:::
+
 :::{tab-item} PyTorch {{ pytorch_icon }}
 
 ```bash
-$ dp --pt freeze -o model.pth
+dp --pt freeze -o model.pth
 ```
 
 in the folder where the model is trained. The output model is called `model.pth`.
@@ -27,16 +40,30 @@ In [multi-task mode](../train/multi-task-training), you need to choose one avail
 to specify which model branch you want to freeze:
 
 ```bash
-$ dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH
+dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH
 ```
 
 The output model is called `model_branch1.pth`, which is the specifically frozen model with the `CHOSEN_BRANCH` head.
 :::
 
+:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}
+
+```bash
+dp --pt-expt freeze -c model.ckpt.pt -o model
+```
+
+The backend writes `.pte` for the dense neighbor-list lower form and `.pt2` for
+the graph lower form. A suffixless output lets DeePMD-kit select the matching
+extension. `--lower-kind graph` requires a graph-eligible model. Conversely, a
+graph-capable DPA model may override a requested `nlist` lower with the graph
+form and emit a warning. In multi-task mode, select a model branch with
+`--head CHOSEN_BRANCH`.
+:::
+
 :::{tab-item} Paddle {{ paddle_icon }}
 
 ```bash
-$ dp --pd freeze -o model
+dp --pd freeze -o model
 ```
 
 in the folder where the model is trained. The output model is called `model.json` and `model.pdiparams`.
@@ -45,10 +72,33 @@ In [multi-task mode](../train/multi-task-training.md), you need to choose one av
 to specify which model branch you want to freeze:
 
 ```bash
-$ dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH
+dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH
 ```
 
 The output model is called `model_branch1.json`, which is the specifically frozen model with the `CHOSEN_BRANCH` head.
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax freeze -c model.ckpt.jax -o model.hlo
+```
+
+The JAX backend can write a StableHLO `.hlo` model, a lossless `.jax` model, or
+a JAX2TF `.savedmodel` model. The `.savedmodel` format requires TensorFlow and
+is the JAX format that supports the C++ inference interface.
+:::
+
 ::::
+
+## Freeze a JAX model with Hessian output {{ jax_icon }}
+
+Use `--hessian` to add coordinate-Hessian output to a frozen JAX energy model:
+
+```bash
+dp --jax freeze -c model.ckpt.jax -o model-hessian.hlo --hessian
+```
+
+The option applies to JAX `.hlo`, `.jax`, and `.savedmodel` outputs. A model
+whose serialized definition already enables Hessian mode retains that mode even
+when `--hessian` is omitted.

doc/model/dpa2.md

@@ -1,7 +1,9 @@
-# Descriptor DPA-2 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }}
+# Descriptor DPA-2 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }}
+> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and
+> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle
+> {{ paddle_icon }}, DP {{ dpmodel_icon }}
 
 The DPA-2 model implementation. See [DPA-2 paper](https://doi.org/10.1038/s41524-024-01493-2) for more details.
 
@@ -96,7 +98,7 @@ The performance improvement will be limited if other parts are more expensive.
 In the pt_expt backend, a graph-eligible DPA-2 descriptor (`repinit/use_three_body` `false` -- the three-body sub-block is not graph-eligible -- and not compressed) can be frozen through a NeighborGraph-native inference path instead of the legacy dense neighbor-list path:
 
 ```bash
-dp --pt_expt freeze -o model.pt2 --lower-kind graph
+dp --pt-expt freeze -o model.pt2 --lower-kind graph
 ```
 
 As with DPA-1's graph path (see [Difference among different backends](train-se-atten.md#difference-among-different-backends)), the graph route considers all neighbors within the cutoff rather than a fixed, padded selection, so its numeric result can differ slightly (down to the AOTInductor floating-point noise floor at non-binding `sel`, larger if `sel` is binding) from the dense/`nlist` path.

doc/model/dpa3.md

@@ -1,7 +1,9 @@
-# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }}
+# Descriptor DPA3 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }}
+> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and
+> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle
+> {{ paddle_icon }}, DP {{ dpmodel_icon }}
 
 DPA3 is an advanced interatomic potential based on message passing.
 As a large atomic model (LAM), it is designed to integrate and jointly train on datasets from different domains,

doc/model/dpa4.md

@@ -1,7 +1,12 @@
-# Descriptor DPA4 {{ pytorch_icon }}
+# Descriptor DPA4 {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: PyTorch {{ pytorch_icon }}
+> **Descriptor backends**: PyTorch and PyTorch-Exportable
+> {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }}
+>
+> The end-to-end `model.type: dpa4` scaffold is supported by PyTorch and
+> PyTorch-Exportable only. JAX support is descriptor-only and does not provide
+> the dedicated DPA4 fitting/model/trainer workflow documented on this page.
 
 DPA4 is the DeePMD-kit implementation of the SeZM (Smooth Equivariant
 Zone-bridging Model) architecture: an SO(3)-equivariant message-passing model
@@ -444,7 +449,7 @@ Two different export routes produce one:
   energy model reports `supports_edge_parallel() == True`, so the archive
   carries the with-comm artifact (`has_comm_artifact=true`) and **supports
   multi-rank LAMMPS out of the box** — no extra freeze options.
-- **pt_expt (`dp --pt_expt freeze`).** Graph-capable models export through the
+- **pt_expt (`dp --pt-expt freeze`).** Graph-capable models export through the
   **NeighborGraph** ABI (see [Graph-native inference route
   (pt_expt)](#graph-native-inference-route-pt_expt) below), which likewise
   embeds a with-comm artifact and supports multi-rank LAMMPS.
@@ -507,7 +512,7 @@ neighbor-list path. Frame-level charge/spin conditioning
 `deepspin` virtual-atom spin scheme remains dense-only:
 
 ```bash
-dp --pt_expt freeze -o model.pt2 --lower-kind graph
+dp --pt-expt freeze -o model.pt2 --lower-kind graph
 ```
 
 As with DPA-1's and DPA-2's graph paths (see [Difference among different
@@ -565,7 +570,7 @@ inference](#multi-gpu-mpi-inference).
 
 - **Native scheme only.** `deepspin`-scheme spin (and the general `spin`
   virtual-atom model outside DPA4/SeZM) is dense-only; only `scheme: native`
-  is graph-eligible. `dp --pt_expt freeze --lower-kind graph` on a
+  is graph-eligible. `dp --pt-expt freeze --lower-kind graph` on a
   `deepspin`-scheme model raises an error at freeze time, per the dense/graph
   eligibility rule above.
 - **Graph route only, no dense fallback.** Unlike a plain-energy DPA4/SeZM
@@ -720,7 +725,9 @@ closed over the one-hop neighbor shell.
 
 ## Limitations
 
-- DPA4/SeZM is implemented for the PyTorch backend only.
+- The end-to-end DPA4/SeZM model and training workflow is implemented for the
+  PyTorch and PyTorch-Exportable backends. JAX provides the descriptor only;
+  the DP implementation is a reference component and does not train models.
 - Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used.
 - Model compression is not supported.
 - Multi-rank (multi-GPU/MPI) LAMMPS inference works for a plain energy model

doc/model/dprc.md

@@ -1,7 +1,9 @@
-# Deep Potential - Range Correction (DPRc) {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
+# Deep Potential - Range Correction (DPRc) {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
+> **Supported backends**: TensorFlow and TensorFlow 2
+> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX
+> {{ jax_icon }}, DP {{ dpmodel_icon }}
 
 Deep Potential - Range Correction (DPRc) is designed to combine with QM/MM method, and corrects energies from a low-level QM/MM method to a high-level QM/MM method:
 

doc/model/linear.md

@@ -1,7 +1,8 @@
 ## Linear model {{ tensorflow_icon }} {{ pytorch_icon }}
 
 > [!NOTE]
-> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}
+> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and
+> PyTorch-Exportable {{ pytorch_icon }}
 
 One can linearly combine existing models with arbitrary coefficients:
 

doc/model/overall.md

@@ -58,7 +58,9 @@ The fitting of the following physical properties is supported
 1. [`ener`](train-energy.md): Fit the energy of the system. The force (derivative with atom positions), the virial (derivative with the box tensor) and the hessian (second-order derivative with atom positions) can also be trained.
 
 > [!WARNING]
-> Due to the restrictions of torch jit script, the models trained with hessian are not jitable so that the frozen models cannot output hessians.
+> The PyTorch TorchScript freeze route cannot output Hessians. The JAX backend
+> can retain Hessian output in a frozen model with
+> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md).
 
 2. [`dipole`](train-fitting-tensor.md): The dipole moment.
 1. [`polar`](train-fitting-tensor.md): The polarizability.