matrix_to_quaternion corner case

Summary: Issue #119. The function `sqrt(max(x, 0))` is not convex and has infinite gradient at 0, but 0 is a subgradient at 0. Here we implement it in such a way as to give 0 as the gradient.

Reviewed By: gkioxari

Differential Revision: D24306294

fbshipit-source-id: 48d136faca083babad4d64970be7ea522dbe9e09

tests/test_rotation_conversions.py

@@ -145,6 +145,20 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
         self.assertEqual(ab.shape, ab_from_matrix.shape)
         self.assertTrue(torch.allclose(ab, ab_from_matrix))
 
+    def test_matrix_to_quaternion_corner_case(self):
+        """Check no bad gradients from sqrt(0)."""
+        matrix = torch.eye(3, requires_grad=True)
+        target = torch.Tensor([0.984808, 0, 0.174, 0])
+
+        optimizer = torch.optim.Adam([matrix], lr=0.05)
+        optimizer.zero_grad()
+        q = matrix_to_quaternion(matrix)
+        loss = torch.sum((q - target) ** 2)
+        loss.backward()
+        optimizer.step()
+
+        self.assertClose(matrix, 0.95 * torch.eye(3))
+
     def test_quaternion_application(self):
         """Applying a quaternion is the same as applying the matrix."""
         quaternions = random_quaternions(3, torch.float64, requires_grad=True)