fix(jax): use more safe_for_vector_norm

njzjz · njzjz · commit 12b5c542cc5a · 2025-06-17T22:31:41.000+08:00
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn> (cherry picked from commit 2433566)
diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
@@ -490,7 +490,7 @@ def call(
         sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))
 
         # get angle nlist (maybe smaller)
-        a_dist_mask = (xp.linalg.vector_norm(diff, axis=-1) < self.a_rcut)[
+        a_dist_mask = (safe_for_vector_norm(diff, axis=-1) < self.a_rcut)[
             :, :, : self.a_sel
         ]
         a_nlist = nlist[:, :, : self.a_sel]

Original file line number	Diff line number	Diff line change
`@@ -490,7 +490,7 @@ def call(`
`490`	`490`	`sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))`
`491`	`491`
`492`	`492`	`# get angle nlist (maybe smaller)`
`493`		`- a_dist_mask = (xp.linalg.vector_norm(diff, axis=-1) < self.a_rcut)[`
	`493`	`+ a_dist_mask = (safe_for_vector_norm(diff, axis=-1) < self.a_rcut)[`
`494`	`494`	`:, :, : self.a_sel`
`495`	`495`	`]`
`496`	`496`	`a_nlist = nlist[:, :, : self.a_sel]`