Increase performance for conversions including axis angles (#1948)

Summary:
This is an extension of https://github.com/facebookresearch/pytorch3d/issues/1544 with various speed, stability, and readability improvements. (I could not find a way to make a commit to the existing PR). This PR is still based on the [Rodrigues' rotation formula](https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Rotation_matrix_%E2%86%94_Euler_axis/angle).

The motivation is the same; this change speeds up the conversions up to 10x, depending on the device, batch size, etc.

### Notes
- As the angles get very close to `π`, the existing implementation and the proposed one start to differ. However, (my understanding is that) this is not a problem as the axis can not be stably inferred from the rotation matrix in this case in general.
- bottler , I tried to follow similar conventions as existing functions to deal with weird angles, let me know if something needs to be changed to merge this.

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

Reviewed By: MichaelRamamonjisoa

Differential Revision: D69193009

Pulled By: bottler

fbshipit-source-id: e5ed34b45b625114ec4419bb89e22a6aefad4eeb

tests/test_rotation_conversions.py

@@ -204,6 +204,9 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
         n_repetitions = 20
         data = torch.rand(n_repetitions, 3)
         matrices = axis_angle_to_matrix(data)
+        self.assertClose(data, matrix_to_axis_angle(matrices), atol=2e-6)
+        self.assertClose(data, matrix_to_axis_angle(matrices, fast=True), atol=2e-6)
+        matrices = axis_angle_to_matrix(data, fast=True)
         mdata = matrix_to_axis_angle(matrices)
         self.assertClose(data, mdata, atol=2e-6)
 
@@ -221,8 +224,10 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
         """mtx -> axis_angle -> mtx"""
         data = random_rotations(13, dtype=torch.float64)
         euler_angles = matrix_to_axis_angle(data)
-        mdata = axis_angle_to_matrix(euler_angles)
-        self.assertClose(data, mdata)
+        euler_angles_fast = matrix_to_axis_angle(data)
+        self.assertClose(data, axis_angle_to_matrix(euler_angles))
+        self.assertClose(data, axis_angle_to_matrix(euler_angles_fast))
+        self.assertClose(data, axis_angle_to_matrix(euler_angles, fast=True))
 
     def test_quaternion_application(self):
         """Applying a quaternion is the same as applying the matrix."""