423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-12-19 13:50:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make some matrix conversion jittable (#898)

Browse Source 
Summary:
Make sure the functions from `rotation_conversion` are jittable, and add some type hints.

Add tests to verify this is the case.

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/898

Reviewed By: patricklabatut

Differential Revision: D31926103

Pulled By: bottler

fbshipit-source-id: bff6013c5ca2d452e37e631bd902f0674d5ca091
This commit is contained in:
una-dinosauria
2021-10-26 14:30:45 -07:00
committed by Facebook GitHub Bot
parent 29417d1f9b
commit bee31c48d3
2 changed files with 61 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/test_rotation_conversions.py
View File

@@ -8,6 +8,7 @@
import itertools
import math
import unittest
from distutils.version import LooseVersion
from typing import Optional, Union
import numpy as np
@@ -264,6 +265,25 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
torch.matmul(r, r.permute(0, 2, 1)), torch.eye(3).expand_as(r), atol=1e-6
)
@unittest.skipIf(LooseVersion(torch.__version__) < "1.9", "recent torchscript only")
def test_scriptable(self):
torch.jit.script(axis_angle_to_matrix)
torch.jit.script(axis_angle_to_quaternion)
torch.jit.script(euler_angles_to_matrix)
torch.jit.script(matrix_to_axis_angle)
torch.jit.script(matrix_to_euler_angles)
torch.jit.script(matrix_to_quaternion)
torch.jit.script(matrix_to_rotation_6d)
torch.jit.script(quaternion_apply)
torch.jit.script(quaternion_multiply)
torch.jit.script(quaternion_to_matrix)
torch.jit.script(quaternion_to_axis_angle)
torch.jit.script(random_quaternions)
torch.jit.script(random_rotation)
torch.jit.script(random_rotations)
torch.jit.script(random_quaternions)
torch.jit.script(rotation_6d_to_matrix)
def _assert_quaternions_close(
self,
input: Union[torch.Tensor, np.ndarray],