Add torch.compile() support for BAE’s PyPose LieTensor operations and sparse Jacobian traversal#41
Conversation
- The dictionary form was unnecessary. Each tensor owns exactly one incoming edge, so id(tensor) → edge was redundant.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces TorchDynamo compatibility for PyPose's LieTensor and enables compiling sparse BSR Jacobians with torch.compile. It refactors the autograd graph traversal to use an immutable optrace structure, implements compile-safe Jacobian component extraction, and updates ambient gradient operations to use pointwise-fusable PyTorch primitives. Feedback on these changes highlights a potential alignment bug in Jacobian materialization when parameters lack components, numerical stability issues in SO3/SE3 exponential maps and Jacobians that could propagate NaN gradients, and opportunities to optimize memory usage by replacing .repeat() with .expand().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces TorchDynamo and torch.compile compatibility for PyPose's LieTensor and BAE's sparse-Jacobian autograd. It implements a compatibility monkeypatch, introduces compile-safe sparse-Jacobian traversal via jacobian_components and materialize_jacobian_components, and replaces complex Lie group operations with plain PyTorch implementations to enable pointwise fusion. Feedback focuses on preventing NaN gradients at zero in the SO3 exponential and Jacobian functions by computing squared norms directly, and optimizing memory efficiency in the graph traversal by replacing tensor repeat operations with zero-copy expand views.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| theta = torch.linalg.norm(input, dim=-1, keepdim=True) | ||
| theta2 = theta.square() | ||
| theta4 = theta2.square() | ||
| nonzero = theta > torch.finfo(theta.dtype).eps | ||
| safe_theta = torch.where(nonzero, theta, torch.ones_like(theta)) | ||
| imaginary_factor = torch.where( | ||
| nonzero, | ||
| torch.sin(0.5 * safe_theta) / safe_theta, | ||
| 0.5 - theta2 / 48.0 + theta4 / 3840.0, | ||
| ) | ||
| real_factor = torch.where( | ||
| nonzero, | ||
| torch.cos(0.5 * theta), | ||
| 1.0 - theta2 / 8.0 + theta4 / 384.0, | ||
| ) |
There was a problem hiding this comment.
In PyTorch, torch.linalg.norm has an undefined (NaN) gradient at zero. When computing theta2 = theta.square(), the chain rule propagates this NaN gradient (2 * theta * d(theta)/d(input) = 0 * NaN = NaN), which then leaks through the else branch of torch.where during the backward pass. Additionally, real_factor uses theta directly in torch.cos(0.5 * theta) instead of safe_theta, which also leads to NaN gradients at zero.
To prevent NaN gradients at zero and improve performance, compute theta2 directly as the sum of squares (input * input).sum(dim=-1, keepdim=True), and use safe_theta in real_factor.
| theta = torch.linalg.norm(input, dim=-1, keepdim=True) | |
| theta2 = theta.square() | |
| theta4 = theta2.square() | |
| nonzero = theta > torch.finfo(theta.dtype).eps | |
| safe_theta = torch.where(nonzero, theta, torch.ones_like(theta)) | |
| imaginary_factor = torch.where( | |
| nonzero, | |
| torch.sin(0.5 * safe_theta) / safe_theta, | |
| 0.5 - theta2 / 48.0 + theta4 / 3840.0, | |
| ) | |
| real_factor = torch.where( | |
| nonzero, | |
| torch.cos(0.5 * theta), | |
| 1.0 - theta2 / 8.0 + theta4 / 384.0, | |
| ) | |
| theta2 = (input * input).sum(dim=-1, keepdim=True) | |
| theta = theta2.sqrt() | |
| theta4 = theta2.square() | |
| nonzero = theta > torch.finfo(theta.dtype).eps | |
| safe_theta = torch.where(nonzero, theta, torch.ones_like(theta)) | |
| imaginary_factor = torch.where( | |
| nonzero, | |
| torch.sin(0.5 * safe_theta) / safe_theta, | |
| 0.5 - theta2 / 48.0 + theta4 / 3840.0, | |
| ) | |
| real_factor = torch.where( | |
| nonzero, | |
| torch.cos(0.5 * safe_theta), | |
| 1.0 - theta2 / 8.0 + theta4 / 384.0, | |
| ) |
| theta = torch.linalg.norm(input, dim=-1, keepdim=True).unsqueeze(-1) | ||
| theta2 = theta.square() |
There was a problem hiding this comment.
Similar to _so3_exp_forward, computing theta2 = theta.square() where theta is the norm introduces NaN gradients at zero due to the chain rule (2 * theta * d(theta)/d(input) = 0 * NaN = NaN). This NaN gradient leaks through the else branch of torch.where for coefficient1 and coefficient2.
To prevent NaN gradients at zero, compute theta2 directly as the sum of squares (input * input).sum(dim=-1, keepdim=True).unsqueeze(-1) and derive theta as theta2.sqrt().
| theta = torch.linalg.norm(input, dim=-1, keepdim=True).unsqueeze(-1) | |
| theta2 = theta.square() | |
| theta2 = (input * input).sum(dim=-1, keepdim=True).unsqueeze(-1) | |
| theta = theta2.sqrt() |
| block_dim = output.shape[-1] | ||
| eye = torch.eye(block_dim, device=output.device, dtype=output.dtype) | ||
| values = eye.unsqueeze(0).repeat(output.shape[0], 1, 1) |
There was a problem hiding this comment.
Using repeat copies the tensor data in memory, which is inefficient. Since values is only used in downstream matrix multiplications or slicing operations, we can use expand to create a zero-copy view instead.
| block_dim = output.shape[-1] | |
| eye = torch.eye(block_dim, device=output.device, dtype=output.dtype) | |
| values = eye.unsqueeze(0).repeat(output.shape[0], 1, 1) | |
| block_dim = output.shape[-1] | |
| eye = torch.eye(block_dim, device=output.device, dtype=output.dtype) | |
| values = eye.unsqueeze(0).expand(output.shape[0], -1, -1) |
| block_dim = output.shape[-1] | ||
| eye = torch.eye( | ||
| block_dim, device=output.device, dtype=output.dtype | ||
| ) | ||
| values = eye.unsqueeze(0).repeat(output.shape[0], 1, 1) |
There was a problem hiding this comment.
Using repeat copies the tensor data in memory, which is inefficient. We can use expand to create a zero-copy view instead.
| block_dim = output.shape[-1] | |
| eye = torch.eye( | |
| block_dim, device=output.device, dtype=output.dtype | |
| ) | |
| values = eye.unsqueeze(0).repeat(output.shape[0], 1, 1) | |
| block_dim = output.shape[-1] | |
| eye = torch.eye( | |
| block_dim, device=output.device, dtype=output.dtype | |
| ) | |
| values = eye.unsqueeze(0).expand(output.shape[0], -1, -1) |
There was a problem hiding this comment.
Code Review
This pull request introduces compatibility improvements for torch.compile with PyPose's LieTensor and ambient gradients, including compile-safe sparse-Jacobian traversal and traceable parameter updates. The review feedback highlights several critical improvements: replacing unsafe torch.Tensor(value) casts with value.view(torch.Tensor) to prevent data copying and device/dtype mismatches; fixing a bug where a tuple slice lie_tensor.shape[-1:] is incorrectly compared to an integer dimension; and preventing potential AttributeErrors by safely retrieving ltype from operands and using getattr for callable names.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _unwrap_tracking_tensor(value): | ||
| if not isinstance(value, TrackingTensor): | ||
| return value | ||
| base = torch.Tensor(value) |
There was a problem hiding this comment.
Using torch.Tensor(value) to cast a subclass tensor to a base torch.Tensor is unsafe in PyTorch. It can trigger data copying, move the tensor to the default device (usually CPU), or reset the dtype to the default (usually float32), which leads to device/dtype mismatches or silent performance degradation.
The standard, safe, and zero-copy way to cast a subclass tensor to a base torch.Tensor while preserving its device, dtype, and autograd history is value.view(torch.Tensor).
| base = torch.Tensor(value) | |
| base = value.view(torch.Tensor) |
There was a problem hiding this comment.
This is to make sure torch dynamo can compile
|
|
||
| def tensor(self) -> torch.Tensor: | ||
| return torch.Tensor.as_subclass(self, torch.Tensor) | ||
| return torch.Tensor(self) |
There was a problem hiding this comment.
Similarly, use self.view(torch.Tensor) instead of torch.Tensor(self) to safely return a base tensor view without copying or risking device/dtype mismatches.
| return torch.Tensor(self) | |
| return self.view(torch.Tensor) |
|
|
||
| def detach(self): | ||
| detached = torch.Tensor.as_subclass(super().detach(), type(self)) | ||
| detached = torch.Tensor(self).detach().as_subclass(type(self)) |
There was a problem hiding this comment.
Use self.view(torch.Tensor) instead of torch.Tensor(self) here as well to safely detach the tensor subclass.
| detached = torch.Tensor(self).detach().as_subclass(type(self)) | |
| detached = self.view(torch.Tensor).detach().as_subclass(type(self)) |
| basics = import_module("pypose.lietensor.basics") | ||
|
|
||
| def _dynamo_compatible_tensor(self): | ||
| return torch.Tensor(self) |
There was a problem hiding this comment.
Use self.view(torch.Tensor) instead of torch.Tensor(self) to safely cast the LieTensor to a base tensor under torch.compile without copying or changing device/dtype.
| return torch.Tensor(self) | |
| return self.view(torch.Tensor) |
| if isinstance(tensor, torch.Tensor) and not isinstance(tensor, cls): | ||
| lie_tensor = torch.Tensor.as_subclass(tensor, lt.LieTensor) | ||
| lie_tensor.ltype = ltype | ||
| if lie_tensor.shape[-1:] != lie_tensor.ltype.dimension: |
There was a problem hiding this comment.
lie_tensor.shape[-1:] returns a torch.Size slice (which is a tuple of length 1, e.g., (7,)), whereas lie_tensor.ltype.dimension is an integer (e.g., 7). Comparing a tuple to an integer (e.g., (7,) != 7) will always evaluate to True.
This causes the "Tensor Shape Invalid" warning to be incorrectly triggered on every single valid operation.
Please use lie_tensor.shape[-1] to get the last dimension as an integer.
| if lie_tensor.shape[-1:] != lie_tensor.ltype.dimension: | |
| if lie_tensor.shape[-1] != lie_tensor.ltype.dimension: |
| # Update the base Tensor instead of assigning through a | ||
| # TrackingTensor/LieTensor view. Mixed-subclass | ||
| # __setitem__ dispatch is not traceable by Dynamo. | ||
| param_tensor = torch.Tensor(param) |
There was a problem hiding this comment.
Use param.view(torch.Tensor) instead of torch.Tensor(param) to avoid copying the parameter to CPU/float32, which would break CUDA execution.
| param_tensor = torch.Tensor(param) | |
| param_tensor = param.view(torch.Tensor) |
| if not pypose_ambient_grad_enabled(): | ||
| return torch.cat([values[..., :6], values[..., 7:]], dim=-1) | ||
| pose = param[..., :7].detach() | ||
| pose = torch.Tensor(param)[..., :7].detach() |
There was a problem hiding this comment.
Use param.view(torch.Tensor) instead of torch.Tensor(param) to safely get the base tensor view without copying or risking device/dtype mismatches.
| pose = torch.Tensor(param)[..., :7].detach() | |
| pose = param.view(torch.Tensor)[..., :7].detach() |
| if isinstance(param, pp.LieTensor): | ||
| if pypose_ambient_grad_enabled(): | ||
| lie_param = param.detach() | ||
| lie_param = torch.Tensor(param).detach() |
There was a problem hiding this comment.
Use param.view(torch.Tensor) instead of torch.Tensor(param) to safely get the base tensor view without copying or risking device/dtype mismatches.
| lie_param = torch.Tensor(param).detach() | |
| lie_param = param.view(torch.Tensor).detach() |
| if func is torch.Tensor.matmul: | ||
| lhs, rhs = args[:2] | ||
| return lhs.ltype.Mul(lhs, rhs) |
There was a problem hiding this comment.
If lhs is a regular Tensor and rhs is a LieTensor (e.g., in tensor @ lie_tensor), lhs will not have an ltype attribute, causing an AttributeError: 'Tensor' object has no attribute 'ltype'.
We should safely retrieve the ltype from either lhs or rhs before calling Mul.
if func is torch.Tensor.matmul:
lhs, rhs = args[:2]
ltype = getattr(lhs, "ltype", None) or getattr(rhs, "ltype", None)
if ltype is not None:
return ltype.Mul(lhs, rhs)|
|
||
| # Access __name__ directly. Dynamo cannot trace hasattr() on the method | ||
| # wrappers used by Tensor properties such as shape. | ||
| if data is not None and func.__name__ in lt.HANDLED_FUNCTIONS: |
There was a problem hiding this comment.
Some callable objects or custom operators (such as OpOverload in PyTorch) might not have a __name__ attribute, which could cause an AttributeError at runtime.
Using getattr(func, "__name__", None) is safer and more robust.
| if data is not None and func.__name__ in lt.HANDLED_FUNCTIONS: | |
| if data is not None and getattr(func, "__name__", None) in lt.HANDLED_FUNCTIONS: |
There was a problem hiding this comment.
Pull request overview
This PR introduces torch.compile() compatibility for PyPose LieTensor operations and BAE’s sparse Jacobian traversal by adding a LieTensor TorchDynamo shim, refactoring optrace to an immutable tuple form, and adding a compile-safe Jacobian “components then materialize” path to work around FakeTensor/AOTAutograd sparse BSR output limitations.
Changes:
- Add a PyPose
LieTensorTorchDynamo compatibility monkeypatch (bae/utils/pypose_compile.py) and ensure it can be enabled via env var / on import. - Refactor
optracerepresentation and add compile-safe sparse Jacobian traversal viajacobian_components()plus eagermaterialize_jacobian_components(). - Add/extend tests covering fullgraph compilation, jacrev correctness, and BAL indexed residual/Jacobian behavior; document activation in README.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/autograd/test_pypose_compile.py | New tests validating fullgraph compile behavior for LieTensor aliasing, ambient-grad ops, jacrev, and LM update paths. |
| tests/autograd/test_pypose_ambient_grad.py | Adds ambient-gradient FD validation for SE(3) Exp. |
| tests/autograd/test_graph_jacobian.py | Extends coverage for tuple-based optrace, component traversal/materialization, and compiled Jacobian comparisons. |
| tests/autograd/test_bal_jacobian.py | Adds a compiled indexed BAL residual + Jacobian-components test; enables ambient-grad patch for BAL tests. |
| README.md | Documents torch.compile enablement and the Jacobian-components materialization workflow. |
| bae/utils/pypose_compile.py | Adds TorchDynamo compatibility shim for PyPose LieTensor (tensor(), __torch_function__, vec2skew). |
| bae/utils/pypose_ambient_grad.py | Makes ambient-grad ops more fusion/compile friendly; adds SO(3)/SE(3) Exp ambient implementations; composes with compile shim. |
| bae/utils/parameter.py | Uses base-tensor aliases (torch.Tensor(param)) to avoid mixed subclass dispatch during trimming. |
| bae/optim/optimizer.py | Updates LM parameter updates to operate through base-tensor aliases for compile compatibility. |
| bae/autograd/graph.py | Introduces component-based Jacobian traversal/materialization and updates trace handling for tuple-based optrace. |
| bae/autograd/function.py | Refactors optrace storage from id(tensor)->edge dict to a direct tuple; introduces compact indexed trace args. |
| bae/init.py | Auto-installs compile patch based on env flag before ambient-grad patch installation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| flat_args, _ = tree_flatten(args) | ||
| ltype = next(arg.ltype for arg in flat_args if isinstance(arg, lt.LieTensor)) |
| for param in params: | ||
| if param.requires_grad: | ||
| numels.append(torch.Size(parameter_update_shape(param)).numel()) | ||
| steps = step.split(numels) | ||
| for (param, d) in zip(params, steps): |
| install_pypose_ambient_grad_monkeypatch() | ||
|
|
||
|
|
`__torch_function__` only flattens `args` when looking for a `LieTensor` to recover `ltype`. If the `LieTensor` is passed via keyword arguments (e.g., `torch.add(input=lie, other=...)`), `args` can be empty and `next(...)` will raise `StopIteration` during tracing/eager execution. Flatten both `args` and `kwargs` when searching for the `LieTensor` source.
Added reference to PR #41 for CUDA graph optimization details.
Summary
This PR adds
torch.compile()support for BAE’s PyPoseLieTensoroperations and sparse Jacobian traversal, including indexed bundle-adjustment parameters.What changed
Adds a compatibility shim for PyPose
LieTensorthat removes TorchDynamo graph breaks caused by:Tensor.as_subclass()when converting aLieTensorto a base tensorhasattr()introspection in__torch_function__requires_gradtensor construction invec2skew@operationsMakes the existing ambient-gradient monkey patch compile-compatible, including SO(3)/SE(3)
Exp,Log,Act,Mul, andInv.Replaces the legacy
id(tensor) → edgeoptracedictionary with a direct tuple representation. Each tracked tensor has exactly one incoming edge, so the dictionary was redundant.Adds compile-safe sparse Jacobian traversal through
jacobian_components():jacrevAdds
materialize_jacobian_components()to construct sparse BSR tensors after the compiled call. PyTorch cannot currently return sparse BSR tensors through FakeTensor/AOTAutograd, so compiled code returns their dense component tensors instead. Single-component BSR materialization is zero-copy.Updates LM parameter retraction and local-Jacobian trimming to operate through base tensor aliases, avoiding mixed tensor-subclass dispatch.
Documents activation through:
export BAE_USE_PYPOSE_TORCH_COMPILE=1The existing
BAE_USE_PYPOSE_AMBIENT_GRAD=1option also installs the compile compatibility shim.Indexing and fusion
Indexed camera and point arguments are stored compactly as
(parent, index)trace metadata and reconstructed inside the compiled graph. This allows Inductor to fuse or otherwise optimize indexed loads with downstream LieTensor and Jacobian operations.fullgraph=Trueguarantees complete graph capture, but it does not guarantee a particular kernel-fusion or intermediate-buffer strategy.Validation
jacrevfor SO(3)/SE(3) operationsIn the tested BA workload, residual and Jacobian-component compilation produced one graph with no graph breaks and reduced runtime from approximately 10.35 ms to 2.21 ms, about a 4.7× speedup.