Further simplify rotation multiplication code

253aff64 · Gustaf Ahdritz · 4e9f7a9d · 253aff64
Commit 253aff64 authored May 07, 2022 by Gustaf Ahdritz
Hide 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)
-    row_2 = row_mul(1)
-    row_3 = row_mul(2)
-    return torch.stack([row_1, row_2, row_3], dim=-2)
+    return torch.stack(
+        [
+            row_mul(0), 
+            row_mul(1), 
+            row_mul(2),
+        ], 
+        dim=-2
+    )


 def rot_vec_mul(