docs(backend): document backend support and workflows

doc/backend.md

@@ -2,10 +2,16 @@
 
 ## Supported backends
 
-DeePMD-kit supports multiple backends: TensorFlow and PyTorch.
+DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2,
+PyTorch-TorchScript, 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-TorchScript 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,25 +21,88 @@ 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.
 
-### PyTorch {{ pytorch_icon }}
+### 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 use a
+`.savedmodeltf` export and therefore require graph-traceable model code.
+Training stores checkpoints in a directory named after the `save_ckpt` prefix
+with `.tf2` appended, such as `model.ckpt.tf2`.
+
+For training, set
+{ref}`training.enable_compile <training/enable_compile>` to `true` to enable
+XLA compilation of the formatted lower-forward path. Setting
+[`DP_JIT`](env.md#envvar-DP_JIT) enables the same model-level default and also
+applies it to SavedModel export. Depending on the workload, compilation may
+improve or reduce performance.
+
+### PyTorch-TorchScript {{ pytorch_icon }}
 
 - Model filename extension: `.pth`
 - Checkpoint filename extension: `.pt`
 
-[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](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.
+
+While `.pth` and `.pt` are the same in the PyTorch package, they have different
+meanings in DeePMD-kit: `.pth` stores a frozen model, while `.pt` stores a
+training 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, change-bias, and testing. Compression support and export
+requirements are documented on the corresponding descriptor pages. 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 PyTorch-TorchScript checkpoint,
+whose parameter names end in `.matrix` and `.bias`.
+
+The `.pt2` suffix identifies an AOTInductor package, but not its lower-input
+ABI. A DPA4/SeZM model frozen with `dp --pt freeze` normally uses the legacy
+`edge_vec` ABI (`lower_input_kind: edge_vec`); its deepspin virtual-atom variant
+uses the dense `nlist` ABI instead. A graph model frozen with
+`dp --pt-expt freeze --lower-kind graph` uses the NeighborGraph ABI
+(`lower_input_kind: graph`). All variants are loaded for inference by the
+PyTorch-Exportable runtime, which reads this metadata to select the correct
+input path. The `--lower-kind` option controls only the PyTorch-Exportable
+freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2)
+for the separate DPA4/SeZM AOTInductor export route.
 
 ### JAX {{ jax_icon }}
 
-- Model filename extension: `.xlo`, `.savedmodel`
-- Checkpoint filename extension: `.jax`
+- DeepEval model filename extensions: `.hlo`, `.savedmodel`
+- Checkpoint and lossless serialization 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. 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.
+
 ### Paddle {{ paddle_icon }}
 
 - Model filename extensions: `.json` and `.pdiparams`
@@ -64,12 +133,19 @@ 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` runtime formats. In particular, `.pt2`
+selects the PyTorch-Exportable inference loader even when the file was produced
+by the DPA4/SeZM `dp --pt freeze` route described above; the archive metadata
+then selects its `edge_vec`, dense `nlist`, or NeighborGraph ABI.
 
 ## Convert model files between backends
 

doc/development/create-a-model-pt.md

@@ -1,9 +1,9 @@
 # Create a model in other backends {{ pytorch_icon }} {{ dpmodel_icon }}
 
 > [!NOTE]
-> **Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
+> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }}, DP {{ dpmodel_icon }}
 >
-> In the following context, we use the PyTorch backend as the example, while it also applies to other backends listed above.
+> In the following context, we use the PyTorch-TorchScript backend as the example, while it also applies to other backends listed above.
 
 If you'd like to create a new model that isn't covered by the existing DeePMD-kit library, but reuse DeePMD-kit's other efficient modules such as data processing, trainer, etc, you may want to read this section.
 
@@ -16,9 +16,9 @@ To incorporate your custom model you'll need to:
 
 ## Design a new component
 
-With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch backend and the framework-independent backend (dpmodel). The PyTorch backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend.
+With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch-TorchScript backend and the framework-independent backend (dpmodel). The PyTorch-TorchScript backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend.
 
-The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch backend.
+The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch-TorchScript backend.
 
 ### New descriptors
 
@@ -118,7 +118,7 @@ class SomeFittingNet(GeneralFitting):
 
 ### New models
 
-The PyTorch backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class.
+The PyTorch-TorchScript backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class.
 
 Subsequently, the `AtomicModel` is encapsulated using the `make_model(AtomicModel)` function, which leverages the `deepmd.pt.model.model.make_model.make_model` function. The purpose of the `make_model` wrapper is to facilitate the translation between atomic property predictions and the extended property predictions and differentiation , e.g. the reduction of atomic energy contribution and the autodiff for calculating the forces and virial. The developers usually need to implement an `AtomicModel` not a `Model`.
 
@@ -217,6 +217,6 @@ When implementing an existing model in a new backend, directly apply the existin
 
 ### Consistent tests
 
-When transferring features from another backend to the PyTorch backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch, and DP (NumPy) through the serialization/deserialization technique.
+When transferring features from another backend to the PyTorch-TorchScript backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch-TorchScript backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch-TorchScript, and DP (NumPy) through the serialization/deserialization technique.
 
-During the development of new components within the PyTorch backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`.
+During the development of new components within the PyTorch-TorchScript backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`.

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-TorchScript and PyTorch-Exportable
+> {{ pytorch_icon }}, JAX {{ jax_icon }}, and DP {{ dpmodel_icon }}.
 
 ## Theory
 
@@ -70,19 +72,68 @@ dp compress -i graph.pb -o graph-compress.pb
 ```
 :::
 
-:::{tab-item} PyTorch {{ pytorch_icon }}
+:::{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. See the
+descriptor documentation for model-specific SavedModel export requirements.
+:::
+
+:::{tab-item} PyTorch-TorchScript {{ 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 dpa1-graph.pt2 -o dpa1-graph-compress.pt2
+```
+
+This command produces an executable compressed artifact only for descriptors
+that support the PyTorch-Exportable compression path. See the
+[DPA-1 model-compression requirements](../model/train-se-atten.md#model-compression)
+for the graph-lowered `.pt2` route used in this example.
+:::
+
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax compress -i frozen_model.jax -o compressed_model.jax
+```
+
+JAX compression accepts `.jax` and `.hlo` inputs. Use `.jax` for the general,
+lossless compressed serialization path. Descriptor pages document whether a
+compressed model can also be exported to StableHLO `.hlo`.
+:::
+
+:::{tab-item} DP {{ dpmodel_icon }}
+
+```bash
+dp --dp compress -i model.dp -o model-compress.dp
+```
+
+DP compression accepts native `.dp`, `.yaml`, and `.yml` 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 compression entrypoints resolve the minimum neighbor distance and tabulate
+supported descriptor embedding networks. 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
@@ -122,7 +173,7 @@ See the documentation of a specific descriptor to see whether it supports model
 
 ## Requirements of installation {{ pytorch_icon }}
 
-When compressing models in the PyTorch backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).
+When compressing models in the PyTorch-TorchScript backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).
 
 The customized OP library for the Python interface is installed by default when building DeePMD-kit from source; see the [installation guide](../install/install-from-source.md) for details.
 

doc/freeze/freeze.md

@@ -8,17 +8,32 @@ 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} PyTorch {{ pytorch_icon }}
+:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}
 
 ```bash
-$ dp --pt freeze -o model.pth
+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`. SavedModel export requires graph-traceable model code;
+descriptor-specific export requirements are documented on the corresponding
+model pages.
+:::
+
+:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }}
+
+```bash
+dp --pt freeze -o model.pth
 ```
 
 in the folder where the model is trained. The output model is called `model.pth`.
@@ -27,16 +42,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 +74,36 @@ 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`
+serialization, or a JAX2TF `.savedmodel` model. The `.hlo` and `.savedmodel`
+formats work with the normal `dp test`/`DeepPot` route; `.jax` is intended for
+checkpoint round-tripping and JAX-MD and is not a DeepEval model format. 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/inference/python.md

@@ -35,7 +35,7 @@ where `descriptors` is the descriptor matrix of the system. This can also be don
 
 > [!NOTE]
 > `eval_descriptor` is the descriptor-only interface supported across backends. In
-> the PyTorch backend, [`eval_embedding`](embedding.md) additionally returns the
+> the PyTorch-TorchScript backend, [`eval_embedding`](embedding.md) additionally returns the
 > descriptor, per-atom feature, and per-structure feature in a single forward pass.
 > PyTorch descriptor/embedding APIs accept `dtype="fp32"`, `"fp64"`, or `"native"`;
 > `eval_descriptor` defaults to `native`, while `eval_embedding` defaults to