make so3_log_map torch script compatible

Summary:
* HAT_INV_SKEW_SYMMETRIC_TOL was a global variable and torch script gives an error when compiling that function. Move it to the function scope.
* torch script gives error when compiling acos_linear_extrapolation because bound is a union of tuple and float. The tuple version is kept in this diff.

Reviewed By: patricklabatut

Differential Revision: D30614916

fbshipit-source-id: 34258d200dc6a09fbf8917cac84ba8a269c00aef

tests/test_so3.py

@@ -7,6 +7,7 @@
 
 import math
 import unittest
+from distutils.version import LooseVersion
 
 import numpy as np
 import torch
@@ -268,6 +269,11 @@ class TestSO3(TestCaseMixin, unittest.TestCase):
                 # all grad values have to be finite
                 self.assertTrue(torch.isfinite(r.grad).all())
 
+    @unittest.skipIf(LooseVersion(torch.__version__) < "1.9", "recent torchscript only")
+    def test_scriptable(self):
+        torch.jit.script(so3_exp_map)
+        torch.jit.script(so3_log_map)
+
     @staticmethod
     def so3_expmap(batch_size: int = 10):
         log_rot = TestSO3.init_log_rot(batch_size=batch_size)