docs(backend): document backend support and workflows

doc/backend.md

@@ -2,10 +2,13 @@
 
 ## 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 {{ tensorflow_icon }}, PyTorch
+{{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP
+{{ dpmodel_icon }} icons are used to mark whether a backend supports a feature.
 
 ### TensorFlow {{ tensorflow_icon }}
 
@@ -15,6 +18,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,6 +40,25 @@ 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`
@@ -64,12 +100,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,8 @@
 # Compress a model {{ tensorflow_icon }} {{ pytorch_icon }}
 
 > [!NOTE]
-> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}
+> **Supported backends**: TensorFlow and TensorFlow 2 {{ tensorflow_icon }},
+> PyTorch and PyTorch-Exportable {{ pytorch_icon }}
 
 ## Theory
 
@@ -70,13 +71,32 @@ 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.
+:::
+
 ::::
 
 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

doc/freeze/freeze.md

@@ -15,6 +15,19 @@ 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 -o model.savedmodeltf
+```
+
+Run the command in the training folder. By default, it reads the
+`model.ckpt.tf2` checkpoint directory and writes the TensorFlow SavedModel to
+`model.savedmodeltf`. Use `-c` to select another checkpoint directory or
+checkpoint prefix. For a multi-task checkpoint, select a branch with
+`--head CHOSEN_BRANCH`.
+:::
+
 :::{tab-item} PyTorch {{ pytorch_icon }}
 
 ```bash
@@ -33,6 +46,19 @@ $ 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. Use `--lower-kind nlist` or `--lower-kind graph` to request a form;
+graph-capable DPA models may select the graph form automatically. In multi-task
+mode, select a model branch with `--head CHOSEN_BRANCH`.
+:::
+
 :::{tab-item} Paddle {{ paddle_icon }}
 
 ```bash

doc/train/training.md

@@ -17,13 +17,27 @@ $ dp --tf train input.json
 ```
 :::
 
+:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}
+
+```bash
+$ dp --tf2 train input.json
+```
+:::
+
 :::{tab-item} PyTorch {{ pytorch_icon }}
 
 ```bash
 $ dp --pt train input.json
 ```
 :::
 
+:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}
+
+```bash
+$ dp --pt-expt train input.json
+```
+:::
+
 :::{tab-item} Paddle {{ paddle_icon }}
 
 ```bash