Skip to content

Preserve PyPose LieTensor semantics in TrackingTensor#20

Merged
zitongzhan merged 4 commits intoreleasefrom
preserveLieTensor
Apr 2, 2026
Merged

Preserve PyPose LieTensor semantics in TrackingTensor#20
zitongzhan merged 4 commits intoreleasefrom
preserveLieTensor

Conversation

@zitongzhan
Copy link
Copy Markdown
Collaborator

Summary

This PR makes TrackingTensor work transparently with PyPose LieTensor inputs so tracked SE(3) parameters keep their LieTensor API instead of degrading to plain tensors.

The main user-facing effect is that code like self.nodes = nn.Parameter(TrackingTensor(nodes)) now preserves pp.SE3 behavior through indexing, so PGO code can use:

  (poses.Inv() @ node1.Inv() @ node2).Log().tensor()

instead of repeatedly re-wrapping everything with pp.SE3(...).

What Changed

  • Made TrackingTensor(...) polymorphic for pp.LieTensor inputs.
  • Added an internal Lie-aware tracking subtype that preserves ltype across:
    • nn.Parameter(...)
    • tensor indexing
    • torch.cat(..., dim=0)
  • Kept the plain tensor tracking path unchanged for non-LieTensor inputs.
  • Updated the Jacobian graph path to preserve LieTensor metadata during torch.vmap(jacrev(...)) via pp.retain_ltype().
  • Kept Jacobian trace shapes unchanged during backward accumulation and only applied trim_SE3_grad when materializing the final Jacobian.
  • Added explicit jactrace cleanup before and after Jacobian construction to avoid stale trace state leaking across calls.
  • Simplified the PGO example to use native LieTensor operations directly.
  • Fixed a test typo in the BAL fixed-camera helper (model.points_3d vs model.points).

Why

Previously, wrapping a pp.LieTensor in TrackingTensor stripped its LieTensor type and ltype, which forced downstream code to re-cast with pp.SE3(...) before every Lie operation.

This PR keeps the sparse-tracing behavior while preserving the LieTensor API, which makes the PGO path cleaner and less error-prone.

  • tracked LieTensor indexing/cat preserving ltype
  • PGO-style SE(3) residuals working without explicit pp.SE3(...) rewraps

…eserving functionality in autograd operations
zitongzhan and others added 2 commits April 1, 2026 20:27
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
@zitongzhan zitongzhan merged commit c9be7bf into release Apr 2, 2026
3 checks passed
@zitongzhan zitongzhan deleted the preserveLieTensor branch April 2, 2026 01:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant