Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting. Args: r: [*, 3, 3] rotation matrices t: [*, 3] coordinate tensors Returns: [*, 3] rotated coordinates
(
r: torch.Tensor,
t: torch.Tensor
)
| 121 | |
| 122 | |
| 123 | def rot_vec_mul( |
| 124 | r: torch.Tensor, |
| 125 | t: torch.Tensor |
| 126 | ) -> torch.Tensor: |
| 127 | """ |
| 128 | Applies a rotation to a vector. Written out by hand to avoid transfer |
| 129 | to avoid AMP downcasting. |
| 130 | |
| 131 | Args: |
| 132 | r: [*, 3, 3] rotation matrices |
| 133 | t: [*, 3] coordinate tensors |
| 134 | Returns: |
| 135 | [*, 3] rotated coordinates |
| 136 | """ |
| 137 | x, y, z = torch.unbind(t, dim=-1) |
| 138 | return torch.stack( |
| 139 | [ |
| 140 | r[..., 0, 0] * x + r[..., 0, 1] * y + r[..., 0, 2] * z, |
| 141 | r[..., 1, 0] * x + r[..., 1, 1] * y + r[..., 1, 2] * z, |
| 142 | r[..., 2, 0] * x + r[..., 2, 1] * y + r[..., 2, 2] * z, |
| 143 | ], |
| 144 | dim=-1, |
| 145 | ) |
| 146 | |
| 147 | |
| 148 | def convert_model_ckpt(state_dict): |
no outgoing calls
no test coverage detected