Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting. Args: a: [*, 3, 3] left multiplicand b: [*, 3, 3] right multiplicand Returns: The product ab
(
a: torch.Tensor,
b: torch.Tensor
)
| 81 | return output_dir |
| 82 | |
| 83 | def rot_matmul( |
| 84 | a: torch.Tensor, |
| 85 | b: torch.Tensor |
| 86 | ) -> torch.Tensor: |
| 87 | """ |
| 88 | Performs matrix multiplication of two rotation matrix tensors. Written |
| 89 | out by hand to avoid AMP downcasting. |
| 90 | |
| 91 | Args: |
| 92 | a: [*, 3, 3] left multiplicand |
| 93 | b: [*, 3, 3] right multiplicand |
| 94 | Returns: |
| 95 | The product ab |
| 96 | """ |
| 97 | def row_mul(i): |
| 98 | return torch.stack( |
| 99 | [ |
| 100 | a[..., i, 0] * b[..., 0, 0] |
| 101 | + a[..., i, 1] * b[..., 1, 0] |
| 102 | + a[..., i, 2] * b[..., 2, 0], |
| 103 | a[..., i, 0] * b[..., 0, 1] |
| 104 | + a[..., i, 1] * b[..., 1, 1] |
| 105 | + a[..., i, 2] * b[..., 2, 1], |
| 106 | a[..., i, 0] * b[..., 0, 2] |
| 107 | + a[..., i, 1] * b[..., 1, 2] |
| 108 | + a[..., i, 2] * b[..., 2, 2], |
| 109 | ], |
| 110 | dim=-1, |
| 111 | ) |
| 112 | |
| 113 | return torch.stack( |
| 114 | [ |
| 115 | row_mul(0), |
| 116 | row_mul(1), |
| 117 | row_mul(2), |
| 118 | ], |
| 119 | dim=-2 |
| 120 | ) |
| 121 | |
| 122 | |
| 123 | def rot_vec_mul( |
no test coverage detected