ci: fix MPS scatter diagnostic

Daniel Zuegner · Daniel Zuegner · commit 9cce4298d71c · 2026-07-27T13:27:26.000Z
Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;

Copilot-Session: 3724d1c0-700a-4e62-956e-5249677cc756
diff --git a/tests/forcefield/test_aoti_compile.py b/tests/forcefield/test_aoti_compile.py
@@ -280,14 +280,15 @@ def run_and_capture(model, input_dict):
     )
     three_basis = model.sbf(rik, torch.acos(cos_jik))
     interaction = model.graph_conv[0].three_body
+    edge_length_column = edge_length.unsqueeze(-1)
     atom_mask = (
         interaction.atom_mlp(atom_attr)[edge_index[0][three_body_indices[:, 1]]]
         * polynomial(
-            edge_length[three_body_indices[:, 0]],
+            edge_length_column[three_body_indices[:, 0]],
             interaction.threebody_cutoff,
         )
         * polynomial(
-            edge_length[three_body_indices[:, 1]],
+            edge_length_column[three_body_indices[:, 1]],
             interaction.threebody_cutoff,
         )
     )

Original file line number	Diff line number	Diff line change
`@@ -280,14 +280,15 @@ def run_and_capture(model, input_dict):`
`280`	`280`	`)`
`281`	`281`	`three_basis = model.sbf(rik, torch.acos(cos_jik))`
`282`	`282`	`interaction = model.graph_conv[0].three_body`
	`283`	`+ edge_length_column = edge_length.unsqueeze(-1)`
`283`	`284`	`atom_mask = (`
`284`	`285`	`interaction.atom_mlp(atom_attr)[edge_index[0][three_body_indices[:, 1]]]`
`285`	`286`	`* polynomial(`
`286`		`- edge_length[three_body_indices[:, 0]],`
	`287`	`+ edge_length_column[three_body_indices[:, 0]],`
`287`	`288`	`interaction.threebody_cutoff,`
`288`	`289`	`)`
`289`	`290`	`* polynomial(`
`290`		`- edge_length[three_body_indices[:, 1]],`
	`291`	`+ edge_length_column[three_body_indices[:, 1]],`
`291`	`292`	`interaction.threebody_cutoff,`
`292`	`293`	`)`
`293`	`294`	`)`