chore: test consistency of rotation matrix

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

Author: njzjz

source/tests/consistent/descriptor/common.py

@@ -74,7 +74,7 @@ def build_tf_descriptor(self, obj, natoms, coords, atype, box, suffix):
         )
         # ensure get_dim_out gives the correct shape
         t_des = tf.reshape(t_des, [1, natoms[0], obj.get_dim_out()])
-        return [t_des], {
+        return [t_des, obj.get_rot_mat()], {
             t_coord: coords,
             t_type: atype,
             t_natoms: natoms,

source/tests/consistent/descriptor/test_dpa1.py

@@ -442,7 +442,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
         )
 
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
-        return (ret[0],)
+        return (ret[0], ret[1])
 
     @property
     def rtol(self) -> float:

source/tests/consistent/descriptor/test_hybrid.py

@@ -168,4 +168,4 @@ def eval_jax(self, jax_obj: Any) -> Any:
         )
 
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
-        return (ret[0],)
+        return (ret[0], ret[1])

source/tests/consistent/descriptor/test_se_atten_v2.py

@@ -340,7 +340,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
         )
 
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
-        return (ret[0],)
+        return (ret[0], ret[1])
 
     @property
     def rtol(self) -> float:

source/tests/consistent/descriptor/test_se_e2_a.py

@@ -259,7 +259,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
         )
 
     def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
-        return (ret[0],)
+        return (ret[0], ret[1])
 
     @property
     def rtol(self) -> float: