Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement SE3.normalize() #506

Merged
merged 12 commits into from
May 2, 2023
Merged

Implement SE3.normalize() #506

Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
55902f1
reimplement normalize for SO3
fantaosha Apr 28, 2023
1a3300b
SO3.normalize() works
fantaosha May 1, 2023
4aa8b22
Change Normalize to use vmap API.
luisenp May 1, 2023
d1b7faa
add a helper to SO3.normalize
fantaosha May 1, 2023
d03f4e9
add tests for `SO3.normalize()`
fantaosha May 1, 2023
2e87575
add `_skew_symm` to simplify the code
fantaosha May 1, 2023
4bd8443
update atol for test_normalize
fantaosha May 1, 2023
47f9c68
add _normalize_backward_helper for reuse
fantaosha May 1, 2023
8f75802
add output types
fantaosha May 1, 2023
1bc60cf
Add SE3.normalize()
fantaosha May 1, 2023
2203a0e
Add SE3.normalize() tests
fantaosha May 1, 2023
990a697
a small refactor of the code
fantaosha May 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
add _skew_symm to simplify the code
  • Loading branch information
@fantaosha
fantaosha committed May 1, 2023
commit 2e875757822f073f0c6ef896ffe2fd46b877a497
9 changes: 5 additions & 4 deletions theseus/labs/lie/functional/so3_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1017,6 +1017,9 @@ def setup_context(ctx, inputs, outputs):

@classmethod
def backward(cls, ctx, grad_output, _):
def _skew_symm(matrix: torch.Tensor) -> torch.Tensor:
return matrix - matrix.transpose(-1, -2)

u, s, v, sign = ctx.saved_tensors
ut = u.transpose(1, 2)
vt = v.transpose(1, 2)
Expand All @@ -1033,13 +1036,11 @@ def backward(cls, ctx, grad_output, _):
F = torch.where(F == 0, grad_output.new_ones(1) * torch.inf, F)
F = F.pow(-1)

u_term: torch.Tensor = u @ (F * (ut @ grad_u - grad_u.transpose(1, 2) @ u))
u_term: torch.Tensor = u @ (F * _skew_symm(ut @ grad_u))
u_term = torch.einsum("n...ij, nj->n...ij", u_term, s)
u_term = u_term @ vt

v_term: torch.Tensor = (
F * (vt @ grad_v - grad_v.transpose(1, 2) @ v)
) @ v.transpose(1, 2)
v_term: torch.Tensor = (F * _skew_symm(vt @ grad_v)) @ v.transpose(1, 2)
v_term = torch.einsum("ni, n...ij->n...ij", s, v_term)
v_term = u @ v_term
return u_term + v_term, None
Expand Down