perf: change order of element-wise op in edge angle update calculations (#4677)

This PR changes the order of element-wise multiply when calculating
`weighted_edge_angle_update`. The largest matrix should be calculated
last to avoid saving large intermediate results and unnecessary
broadcast.

I've tested this PR on OMat with 9 DPA-3 layers and batch size=auto:512.
| Metric                 | Before   | After    | Improvement |
|------------------------|----------|----------|-------------|
| Peak Memory  | 25.0G    | 21.4G    | -15%         |
| Speed (per 100 steps)  | 31.27s   | 28.9s    | +7.7%       |

Since this is an element-wise multiply, changing the order of arguments
should not affect the result. The correctness is verified by
`torch.allclose`.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Refactor**
- Improved internal calculation order for weighted updates to enhance
code clarity and maintainability, while ensuring the functionality
remains consistent.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: caic99

deepmd/dpmodel/descriptor/repflows.py

@@ -1117,9 +1117,9 @@ def call(
 
             # nb x nloc x a_nnei x a_nnei x e_dim
             weighted_edge_angle_update = (
-                edge_angle_update
-                * a_sw[:, :, :, xp.newaxis, xp.newaxis]
+                a_sw[:, :, :, xp.newaxis, xp.newaxis]
                 * a_sw[:, :, xp.newaxis, :, xp.newaxis]
+                * edge_angle_update
             )
             # nb x nloc x a_nnei x e_dim
             reduced_edge_angle_update = xp.sum(weighted_edge_angle_update, axis=-2) / (

deepmd/pt/model/descriptor/repflow_layer.py

@@ -698,9 +698,9 @@ def forward(
 
             # nb x nloc x a_nnei x a_nnei x e_dim
             weighted_edge_angle_update = (
-                edge_angle_update
-                * a_sw[:, :, :, None, None]
+                a_sw[:, :, :, None, None]
                 * a_sw[:, :, None, :, None]
+                * edge_angle_update
             )
             # nb x nloc x a_nnei x e_dim
             reduced_edge_angle_update = torch.sum(