docs(jax): document training and freeze --hessian

doc/backend.md

@@ -25,15 +25,18 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different
 
 ### JAX {{ jax_icon }}
 
-- Model filename extension: `.xlo`, `.savedmodel`
+- Model filename extension: `.hlo`, `.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`

doc/freeze/freeze.md

@@ -51,4 +51,27 @@ $ 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/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.

doc/model/train-energy-hessian.md

@@ -91,9 +91,12 @@ The detailed loss can be found in `lcurve.out`:
 ## Test the Model
 
 > [!WARNING]
-> A model trained with Hessian cannot be frozen. If freezing is enforced, the model will be treated as a standard energy model, and the frozen one will no longer be able to output Hessian predictions.
+> The PyTorch TorchScript freeze route does not preserve Hessian output. A
+> PyTorch model frozen with `dp --pt freeze` is treated as a standard energy
+> model. The JAX backend can preserve Hessian output in a frozen model with
+> `dp --jax freeze --hessian`.
 
-If one do freeze and test a Hessian model using the commands:
+If one freezes and tests a Hessian model through the PyTorch route:
 
 ::::{tab-set}
 
@@ -116,7 +119,8 @@ ${output_prefix}.e.out   ${output_prefix}.e_peratom.out  ${output_prefix}.f.out
 ${output_prefix}.v.out   ${output_prefix}.v_peratom.out
 ```
 
-If one intends to use the trained model for Hessian predictions, then he/she is supposed to test the model directly without performing a freezing operation:
+For PyTorch Hessian predictions, test the training checkpoint directly without
+freezing:
 
 ::::{tab-set}
 

doc/train/training.md

@@ -39,6 +39,17 @@ $ CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json
 ```
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+$ dp --jax train input.json
+```
+
+JAX training, introduced after DeePMD-kit v3.1.3, writes checkpoints with the
+`.jax` extension. It supports fine-tuning and multi-task configurations, but
+does not currently support `--init-frz-model`.
+:::
+
 ::::
 
 where `input.json` is the name of the input script.