fix(jax): fix several serialization and jit issues

- `deepmd/jax/descriptor/__init__.py` imports SeT and DPA-2 to let them found by the plugin; - `deepmd/dpmodel/descriptor/dpa1.py` fixes the jit issue regarding to the shape generated by `jnp.prod`. The shape should be static by using `math.prod`. - `deepmd/jax/model/ener_model.py` and `deepmd/jax/model/dp_zbl_model.py` stop the graident of coordinates when rebuilding the neighbor list. The graient of sort causes an error due to jax-ml/jax#24730. Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 5, 2024 · 54dc410 · 54dc410
1 parent dabedd2
commit 54dc410
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 2 deletions.
diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import math
 from typing import (
     Any,
     Callable,
@@ -852,7 +853,7 @@ def cal_g(
     ):
         xp = array_api_compat.array_namespace(ss)
         nfnl, nnei = ss.shape[0:2]
-        shape2 = xp.prod(xp.asarray(ss.shape[2:]))
+        shape2 = math.prod(ss.shape[2:])
         ss = xp.reshape(ss, (nfnl, nnei, shape2))
         # nfnl x nnei x ng
         gg = self.embeddings[embedding_idx].call(ss)
@@ -866,7 +867,7 @@ def cal_g_strip(
         assert self.embeddings_strip is not None
         xp = array_api_compat.array_namespace(ss)
         nfnl, nnei = ss.shape[0:2]
-        shape2 = xp.prod(xp.asarray(ss.shape[2:]))
+        shape2 = math.prod(ss.shape[2:])
         ss = xp.reshape(ss, (nfnl, nnei, shape2))
         # nfnl x nnei x ng
         gg = self.embeddings_strip[embedding_idx].call(ss)

diff --git a/deepmd/jax/descriptor/__init__.py b/deepmd/jax/descriptor/__init__.py
@@ -2,6 +2,9 @@
 from deepmd.jax.descriptor.dpa1 import (
     DescrptDPA1,
 )
+from deepmd.jax.descriptor.dpa2 import (
+    DescrptDPA2,
+)
 from deepmd.jax.descriptor.hybrid import (
     DescrptHybrid,
 )
@@ -14,11 +17,16 @@
 from deepmd.jax.descriptor.se_t import (
     DescrptSeT,
 )
+from deepmd.jax.descriptor.se_t_tebd import (
+    DescrptSeTTebd,
+)
 
 __all__ = [
     "DescrptSeA",
     "DescrptSeR",
     "DescrptSeT",
+    "DescrptSeTTebd",
     "DescrptDPA1",
+    "DescrptDPA2",
     "DescrptHybrid",
 ]
diff --git a/deepmd/jax/model/dp_zbl_model.py b/deepmd/jax/model/dp_zbl_model.py
@@ -12,6 +12,7 @@
     flax_module,
 )
 from deepmd.jax.env import (
+    jax,
     jnp,
 )
 from deepmd.jax.model.base_model import (
@@ -48,3 +49,18 @@ def forward_common_atomic(
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
         )
+
+    def format_nlist(
+        self,
+        extended_coord: jnp.ndarray,
+        extended_atype: jnp.ndarray,
+        nlist: jnp.ndarray,
+        extra_nlist_sort: bool = False,
+    ):
+        return DPZBLModelDP.format_nlist(
+            self,
+            jax.lax.stop_gradient(extended_coord),
+            extended_atype,
+            nlist,
+            extra_nlist_sort=extra_nlist_sort,
+        )
diff --git a/deepmd/jax/model/ener_model.py b/deepmd/jax/model/ener_model.py
@@ -12,6 +12,7 @@
     flax_module,
 )
 from deepmd.jax.env import (
+    jax,
     jnp,
 )
 from deepmd.jax.model.base_model import (
@@ -48,3 +49,18 @@ def forward_common_atomic(
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
         )
+
+    def format_nlist(
+        self,
+        extended_coord: jnp.ndarray,
+        extended_atype: jnp.ndarray,
+        nlist: jnp.ndarray,
+        extra_nlist_sort: bool = False,
+    ):
+        return EnergyModelDP.format_nlist(
+            self,
+            jax.lax.stop_gradient(extended_coord),
+            extended_atype,
+            nlist,
+            extra_nlist_sort=extra_nlist_sort,
+        )