From e9f9321b30f9d492f5fc387cbe9e7beddab8a0fc Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 21 Nov 2024 05:37:39 -0500
Subject: [PATCH] docs: document more differences among different backends
 (#4388)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced new sections titled "Difference among different backends"
in multiple documents, clarifying backend-specific constraints for
various descriptors.
- Added new descriptors for model training, including `"se_atten"` and
`"se_atten_v2"`.

- **Documentation**
- Enhanced clarity and structure of existing documents regarding tensor
fitting and model training.
- Updated guidance on model compression options and backend limitations.

- **Bug Fixes**
- Minor grammatical and formatting adjustments made for improved
readability.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 doc/model/train-fitting-tensor.md | 8 +++++++-
 doc/model/train-se-atten.md       | 7 ++++++-
 doc/model/train-se-e2-a.md        | 5 +++++
 doc/model/train-se-e2-r.md        | 5 +++++
 doc/model/train-se-e3.md          | 4 ++++
 5 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/doc/model/train-fitting-tensor.md b/doc/model/train-fitting-tensor.md
index d4d546eccf..1463bc04b7 100644
--- a/doc/model/train-fitting-tensor.md
+++ b/doc/model/train-fitting-tensor.md
@@ -30,7 +30,8 @@ $deepmd_source_dir/examples/water_tensor/polar/polar_input_torch.json
 
 The training and validation data are also provided our examples. But note that **the data provided along with the examples are of limited amount, and should not be used to train a production model.**
 
-Similar to the `input.json` used in `ener` mode, training JSON is also divided into {ref}`model <model>`, {ref}`learning_rate <learning_rate>`, {ref}`loss <loss>` and {ref}`training <training>`. Most keywords remain the same as `ener` mode, and their meaning can be found [here](train-se-e2-a.md). To fit a tensor, one needs to modify {ref}`model[standard]/fitting_net <model[standard]/fitting_net>` and {ref}`loss <loss>`.
+Similar to the `input.json` used in `ener` mode, training JSON is also divided into {ref}`model <model>`, {ref}`learning_rate <learning_rate>`, {ref}`loss <loss>` and {ref}`training <training>`. Most keywords remain the same as `ener` mode, and their meaning can be found [here](train-se-e2-a.md).
+To fit a tensor, one needs to modify {ref}`fitting_net <model[standard]/fitting_net>` and {ref}`loss <loss>`.
 
 ## Theory
 
@@ -241,3 +242,8 @@ One may notice that in each step, some of the local loss and global loss will be
 ```
 
 During training, at each step when the `lcurve.out` is printed, the system used for evaluating the training (validation) error may be either with only global or only atomic labels, thus the corresponding atomic or global errors are missing and are printed as zeros.
+
+## Difference among different backends
+
+To only fit against a subset of atomic types, in the TensorFlow backend, {ref}`fitting_net/sel_type <model[standard]/fitting_net[dipole]/sel_type>` should be set to selected types;
+in other backends, {ref}`atom_exclude_types <model/atom_exclude_types>` should be set to excluded types.
diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md
index d8d7c3c98c..92a56395f6 100644
--- a/doc/model/train-se-atten.md
+++ b/doc/model/train-se-atten.md
@@ -152,11 +152,16 @@ In the TensorFlow backend, the {ref}`type_embedding <model/type_embedding>` sect
 
 In other backends, type embedding is within this descriptor with the {ref}`tebd_dim <model[standard]/descriptor[se_atten_v2]/tebd_dim>` argument.
 
-## Difference between TensorFlow and other backends
+## Difference among different backends
 
 TensorFlow and other backends have different implementations for {ref}`smooth_type_embedding <model[standard]/descriptor[se_atten_v2]/smooth_type_embedding>`.
 The results are inconsistent when `smooth_type_embedding` is `true`.
 
+In the TensorFlow backend, {ref}`scaling_factor <model[standard]/descriptor[se_atten]/scaling_factor>` cannot set to a value other than `1.0`;
+{ref}`normalize <model[standard]/descriptor[se_atten]/normalize>` cannot be set to `false`;
+{ref}`temperature <model[standard]/descriptor[se_atten]/temperature>` cannot be set;
+{ref}`concat_output_tebd <model[standard]/descriptor[se_atten]/concat_output_tebd>` cannot be set to `false`.
+
 ## Type map
 
 For training large systems, especially those with dozens of elements, the {ref}`type <model/type_map>` determines the element index of training data:
diff --git a/doc/model/train-se-e2-a.md b/doc/model/train-se-e2-a.md
index 2cc537e349..5143d5b5fb 100644
--- a/doc/model/train-se-e2-a.md
+++ b/doc/model/train-se-e2-a.md
@@ -100,6 +100,11 @@ The construction of the descriptor is given by section {ref}`descriptor <model[s
 Type embdding is only supported in the TensorFlow backends.
 `se_e2_a` with type embedding and [`se_atten`](./train-se-atten.md) (or its updated version) without any attention layer are mathematically equivalent, so `se_atten` can be a substitute in other backends.
 
+## Difference among different backends
+
+In the TensorFlow backend, {ref}`env_protection <model[standard]/descriptor[se_e2_a]/env_protection>` cannot be set to a non-zero value.
+In the JAX backend, {ref}`type_one_side <model[standard]/descriptor[se_e2_a]/type_one_side>` cannot be set to `false`.
+
 ## Model compression
 
 Model compression is supported when type embedding is not used.
diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md
index dbaf8bf364..62c7311a08 100644
--- a/doc/model/train-se-e2-r.md
+++ b/doc/model/train-se-e2-r.md
@@ -74,6 +74,11 @@ The type of the descriptor is set by the key {ref}`type <model[standard]/descrip
 
 Type embdding is only supported in the TensorFlow backends.
 
+## Difference among different backends
+
+In the TensorFlow backend, {ref}`env_protection <model[standard]/descriptor[se_e2_r]/env_protection>` cannot be set to a non-zero value.
+In the PyTorch, JAX, and DP backend, {ref}`type_one_side <model[standard]/descriptor[se_e2_r]/type_one_side>` cannot be set to `false`.
+
 ## Model compression
 
 Model compression is supported when type embedding is not used.
diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md
index 4e7f4fa53a..2f1b3ba972 100644
--- a/doc/model/train-se-e3.md
+++ b/doc/model/train-se-e3.md
@@ -69,6 +69,10 @@ The type of the descriptor is set by the key {ref}`type <model[standard]/descrip
 
 Use [`se_e3_tebd`](./train-se-e3-tebd.md) for type embedding support.
 
+## Difference among different backends
+
+In the TensorFlow backend, {ref}`env_protection <model[standard]/descriptor[se_e3]/env_protection>` cannot be set to a non-zero value.
+
 ## Model compression
 
 Model compression is supported.