Further simplify rotation multiplication code

253aff64 · Gustaf Ahdritz · 4e9f7a9d · 253aff64
Commit 253aff64 authored May 07, 2022 by Gustaf Ahdritz
Show whitespace changes
Inline Side-by-side

Showing with 8 additions and 4 deletions

openfold/utils/rigid_utils.py openfold/utils/rigid_utils.py +8 -4

No files found.
--- a/openfold/utils/rigid_utils.py
+++ b/openfold/utils/rigid_utils.py
@@ -50,10 +50,14 @@ def rot_matmul(
            dim=-1,
        )
-    row_1 = row_mul(0)
+    return torch.stack(
-    row_2 = row_mul(1)
+        [
-    row_3 = row_mul(2)
+            row_mul(0), 
-    return torch.stack([row_1, row_2, row_3], dim=-2)
+            row_mul(1), 
+            row_mul(2),
+        ], 
+        dim=-2
+    )
 def rot_vec_mul(