$ show pypose/bae
bae is a PyTorch library for exact second-order optimization, specializing in Bundle Adjustment (BA) and Pose Graph Optimization (PGO) for robotics and 3D vision. It provides a sparse autograd system that automatically computes sparse block Jacobians by tracing operations on tracked parameters, enabling efficient Levenberg-Marquardt and Schur complement optimizers. The library includes CUDA-accelerated sparse matrix operations via cuSPARSE and Warp, ambient gradient implementations for SE(3) operations that are torch.compile compatible, and utilities for loading BAL and PGO datasets. It is designed as a research library with clean abstractions for defining residual graphs and optimizing over Lie groups. The architecture separates autograd tracing (bae/autograd), optimization algorithms (bae/optim), sparse matrix utilities (bae/sparse), and supporting utilities (bae/utils). The sparse autograd system uses a custom Tensor subclass (TrackingTensor) that records index, map, and cat operations via __torch_function__, then computes sparse BSR Jacobians via a backward traversal. The LM and Schur optimizers support trust region and adaptive strategies, with PCG as the default linear solver. The library also provides monkeypatches for PyPose LieTensor to enable torch.compile full-graph capture, and ambient (plain PyTorch) implementations of SE(3) operations for gradient compatibility.
The LLM identified 6 files worth highlighting for reuse:
Provides analytical Jacobian of SE(3) retraction (exponential map) with respect to ambient coordinates. Used to transform Jacobians from tangent space to ambient space for parameter updates. Self-contained, only depends on torch, and demonstrates non-trivial Lie group calculus.
Implements a matrix-free normal equation operator that represents J^T J x + damping * diag * x without materializing the full J^T J matrix. Supports dense, CSR, COO, and BSR Jacobian layouts. Reusable for any iterative solver (e.g., CG) that needs to apply the normal operator without explicit construction.
Replaces PyPose's custom autograd for SE(3) operations (Act, Mul, Inv, Log, Exp) with pure PyTorch implementations that are torch.compile compatible. Includes analytical forward and gradient computations for Lie group operations. Demonstrates significant skill in Lie group mathematics and PyTorch internals.
Provides conversion functions between PyTorch sparse BSR tensors and Warp's BSR matrix format, enabling use of Warp's optimized sparse operations (bsr_mm, bsr_mv, bsr_axpy) on PyTorch tensors. Also registers custom CUDA implementations for torch.add on BSR tensors. Reusable pattern for interoperability between PyTorch and Warp.
Monkeypatches PyPose's LieTensor to be traceable by TorchDynamo, enabling torch.compile with LieTensors. Fixes issues with as_subclass, hasattr, and vec2skew that otherwise break graph capture. Essential for compiling BA/PGO models. Self-contained and demonstrates deep knowledge of PyTorch internals.
Defines TrackingTensor, a torch.Tensor subclass that intercepts operations via __torch_function__ to record a trace of index, map, and cat operations. This trace is later used by the sparse autograd system to compute sparse Jacobians. The design enables automatic sparse Jacobian computation without user intervention, a non-trivial achievement.