Migrate SO3 (reissue) (#539)

* Migrate core SO3 functions to use torchlie. * Replaced rotate and some tensor checks. * Fix failing unit test. * Removed debug statement.
facebookresearch · Jun 21, 2023 · 18a46a0 · 18a46a0
1 parent 2f77df7
commit 18a46a0
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 402 deletions.
diff --git a/tests/theseus_tests/core/test_vectorizer.py b/tests/theseus_tests/core/test_vectorizer.py
@@ -206,7 +206,7 @@ def _check_vectorized_wrappers(vectorization, objective):
         for cost_fn in objective.cost_functions.values():
             if cost_fn is w.cost_fn:
                 w_jac, w_err = cost_fn.weighted_jacobians_error()
-                assert w._cached_error.allclose(w_err)
+                torch.testing.assert_close(w._cached_error, w_err)
                 for jac, exp_jac in zip(w._cached_jacobians, w_jac):
                     torch.testing.assert_close(jac, exp_jac, atol=1e-6, rtol=1e-6)
 

diff --git a/theseus/geometry/se3.py b/theseus/geometry/se3.py
@@ -583,7 +583,7 @@ def transform_from(
             # Right jacobians for SE3 are computed
             Jg = torch.zeros(batch_size, 3, 6, dtype=self.dtype, device=self.device)
             Jg[:, :, :3] = self[:, :, :3]
-            Jg[:, :, 3:] = -self[:, :, :3] @ SO3.hat(p)
+            Jg[:, :, 3:] = -self[:, :, :3] @ SO3.hat(p.squeeze(-1))
             # Jacobians for point
             Jpnt = Jg[:, :, :3]