resolve comments

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Oct 25, 2024 · b9eefd3 · b9eefd3
1 parent 004b89a
commit b9eefd3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 20 deletions.
diff --git a/deepmd/dpmodel/model/transform_output.py b/deepmd/dpmodel/model/transform_output.py
@@ -102,18 +102,16 @@ def communicate_extended_output(
                     )
                     # jax only
                     if array_api_compat.is_jax_array(force):
-                        from deepmd.jax.env import (
-                            jnp,
+                        from deepmd.jax.common import (
+                            scatter_sum,
                         )
 
-                        f_idx = xp.arange(force.size, dtype=xp.int64).reshape(
-                            force.shape
+                        force = scatter_sum(
+                            force,
+                            1,
+                            mapping,
+                            model_ret[kk_derv_r],
                         )
-                        new_idx = jnp.take_along_axis(f_idx, mapping, axis=1).ravel()
-                        f_shape = force.shape
-                        force = force.ravel()
-                        force = force.at[new_idx].add(model_ret[kk_derv_r].ravel())
-                        force = force.reshape(f_shape)
                     else:
                         raise NotImplementedError("Only JAX arrays are supported.")
                     new_ret[kk_derv_r] = force
@@ -132,18 +130,16 @@ def communicate_extended_output(
                     )
                     # jax only
                     if array_api_compat.is_jax_array(virial):
-                        from deepmd.jax.env import (
-                            jnp,
+                        from deepmd.jax.common import (
+                            scatter_sum,
                         )
 
-                        v_idx = xp.arange(virial.size, dtype=xp.int64).reshape(
-                            virial.shape
+                        virial = scatter_sum(
+                            virial,
+                            1,
+                            mapping,
+                            model_ret[kk_derv_c],
                         )
-                        new_idx = jnp.take_along_axis(v_idx, mapping, axis=1).ravel()
-                        v_shape = virial.shape
-                        virial = virial.ravel()
-                        virial = virial.at[new_idx].add(model_ret[kk_derv_c].ravel())
-                        virial = virial.reshape(v_shape)
                     else:
                         raise NotImplementedError("Only JAX arrays are supported.")
                     new_ret[kk_derv_c] = virial

diff --git a/deepmd/jax/common.py b/deepmd/jax/common.py
@@ -95,3 +95,13 @@ def __dlpack__(self, *args, **kwargs):
 
     def __dlpack_device__(self, *args, **kwargs):
         return self.value.__dlpack_device__(*args, **kwargs)
+
+
+def scatter_sum(input, dim, index: jnp.ndarray, src: jnp.ndarray) -> jnp.ndarray:
+    """Reduces all values from the src tensor to the indices specified in the index tensor."""
+    idx = jnp.arange(input.size, dtype=jnp.int64).reshape(input.shape)
+    new_idx = jnp.take_along_axis(idx, index, axis=dim).ravel()
+    shape = input.shape
+    input = input.ravel()
+    input = input.at[new_idx].add(src.ravel())
+    return input.reshape(shape)
diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py
@@ -57,7 +57,16 @@ def forward_common_atomic(
                 for ss in range(size):
 
                     def eval_output(
-                        cc_ext, extended_atype, nlist, mapping, fparam, aparam
+                        cc_ext,
+                        extended_atype,
+                        nlist,
+                        mapping,
+                        fparam,
+                        aparam,
+                        *,
+                        _kk=kk,
+                        _ss=ss,
+                        _atom_axis=atom_axis,
                     ):
                         atomic_ret = self.atomic_model.forward_common_atomic(
                             cc_ext[None, ...],
@@ -67,7 +76,7 @@ def eval_output(
                             fparam=fparam[None, ...] if fparam is not None else None,
                             aparam=aparam[None, ...] if aparam is not None else None,
                         )
-                        return jnp.sum(atomic_ret[kk][0], axis=atom_axis)[ss]
+                        return jnp.sum(atomic_ret[_kk][0], axis=_atom_axis)[_ss]
 
                     ffi = -jax.vmap(jax.grad(eval_output, argnums=0))(
                         extended_coord,