Add torch.compile() support for BAE’s PyPose LieTensor operations and sparse Jacobian traversal

README.md

@@ -82,11 +82,10 @@
 - **Schur Complement Optimizer**: Solve a Schur-reduce linear system for [optimized memory consumption](https://github.com/pypose/bae/blob/release/doc/memory_performance.md) 
 
 ### Future Plan
-- [ ] Add Apple Silicon GPU support, [PyTorch PR WIP](https://github.com/pytorch/pytorch/pull/177757)
-- [ ] Reduce runtime overhead using CUDA graph
+- [ ] Reduce runtime overhead using CUDA graph (WIP in dynamo branch: compile fwd & backward with `torch.compile()` reducing latency from 10ms to 2.2ms)
 - [ ] Distributed Tensor (DTensor) and FSDP support for multi-GPU and distributed optimization
-- [ ] An new backend for [distributed solver](https://github.com/NVIDIA/AMGX)
 - [x] Schur complement (added in [PR #35](https://github.com/pypose/bae/pull/35))
+- [ ] Add Apple Silicon GPU support, [PyTorch PR WIP](https://github.com/pytorch/pytorch/pull/177757)
 
 ## Installation
 
@@ -133,6 +132,66 @@ python -m pip install git+https://github.com/pypose/bae.git
    python -m pip install --no-build-isolation -v -e .  # following https://github.com/pytorch/pytorch
    ```
 
+### `torch.compile` with `LieTensor`
+
+Enable the `LieTensor` TorchDynamo compatibility shim before importing either
+`pypose` or `bae`:
+
+```bash
+export BAE_USE_PYPOSE_TORCH_COMPILE=1
+```
+
+The existing ambient-gradient implementation automatically enables the same
+shim, so `BAE_USE_PYPOSE_AMBIENT_GRAD=1` is sufficient when ambient gradients
+are required. The shim can also be installed explicitly before compiling a
+model or residual function:
+
+```python
+import torch
+
+from bae.utils.pypose_compile import install_pypose_torch_compile_monkeypatch
+
+install_pypose_torch_compile_monkeypatch()
+compiled_model = torch.compile(model, fullgraph=True)
+```
+
+With `fullgraph=True`, indexed `sjac=True` parameters retain their sparse
+Jacobian dependency trace without forcing the gathered camera and point blocks
+to escape the compiled graph. This gives Inductor the opportunity to load
+permuted rows directly inside fused kernels. Inductor will decide the most efficient way, 
+whether to materialize standalone indexed tensors when the
+gathered values have multiple downstream consumers, as can happen during
+Jacobian computation. `fullgraph=True` guarantees graph capture but does not ensure 
+particular kernel-fusion or buffer-allocation strategy.
+
+PyTorch cannot currently represent a sparse BSR tensor as a FakeTensor/AOT
+graph output. To compile the residual and sparse-Jacobian traversal together,
+return its dense component tensors from the compiled function and materialize
+the BSR wrapper immediately afterward:
+
+```python
+from bae.autograd.graph import (
+    jacobian_components,
+    materialize_jacobian_components,
+)
+
+def residual_and_jacobian(observations, camera_indices, point_indices):
+    residual = model(observations, camera_indices, point_indices)
+    components = jacobian_components(
+        residual, (model.pose, model.points)
+    )
+    return residual, components
+
+compiled = torch.compile(residual_and_jacobian, fullgraph=True)
+residual, components = compiled(observations, camera_indices, point_indices)
+jacobians = materialize_jacobian_components(components)
+```
+
+The component traversal and Jacobian values are compiled. For one component per
+parameter, as in the BA residual above, BSR materialization is an eager,
+zero-copy wrapper operation. Graphs with multiple contributions to the same
+parameter additionally combine those sparse components after materialization.
+
 ## Agent Skills
 
 This repo includes skills in [.agent/skills](.agent/skills):

bae/__init__.py

@@ -1,3 +1,5 @@
+from .utils.pypose_compile import maybe_install_pypose_torch_compile_monkeypatch
 from .utils.pypose_ambient_grad import maybe_install_pypose_ambient_grad_monkeypatch
 
+maybe_install_pypose_torch_compile_monkeypatch()
 maybe_install_pypose_ambient_grad_monkeypatch()

bae/autograd/function.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pypose as pp
 from functools import wraps
+from torch.utils._pytree import tree_map
 
 WHITELISTED_MAPS = tuple(
     func for func in (
@@ -23,6 +24,8 @@
     torch._C.TensorBase.to,
 }
 
+_INDEXED_TRACE_TAG = "bae.indexed"
+
 
 def _iter_tracked_tensors(values):
     if isinstance(values, torch.Tensor):
@@ -36,9 +39,7 @@ def _iter_tracked_tensors(values):
 
 
 def _attach_index_trace(result, index, tensor):
-    if not hasattr(result, 'optrace'):
-        result.optrace = {}
-    result.optrace[id(result)] = ("index", index, tensor)
+    result.optrace = ("index", index, tensor)
     return result
 
 
@@ -47,18 +48,31 @@ def _attach_cat_trace(result, tensors, dim):
     offset = 0
     for t in tensors:
         n = t.shape[0]
-        if hasattr(t, 'optrace') or isinstance(t, torch.nn.Parameter):
+        if (
+            hasattr(t, "optrace")
+            or isinstance(t, TrackingTensor)
+            or isinstance(t, torch.nn.Parameter)
+        ):
             tracked.append((offset, offset + n, t))
         offset += n
-    result.optrace = {id(result): ("cat", dim, tuple(tracked))}
+    result.optrace = ("cat", dim, tuple(tracked))
     return result
 
 
 def _attach_map_trace(result, func, args):
-    result.optrace = {id(result): ("map", func, args)}
+    compact_args = tuple(_compact_map_arg(arg) for arg in args)
+    result.optrace = ("map", func, compact_args)
     return result
 
 
+def _compact_map_arg(arg):
+    if isinstance(arg, TrackingTensor) and hasattr(arg, "optrace"):
+        trace = arg.optrace
+        if trace[0] == "index":
+            return (_INDEXED_TRACE_TAG, trace[2], trace[1])
+    return arg
+
+
 def _find_tracking_source(values, cls):
     for value in _iter_tracked_tensors(values):
         if isinstance(value, cls):
@@ -69,23 +83,16 @@ def _find_tracking_source(values, cls):
 def _unwrap_tracking_tensor(value):
     if not isinstance(value, TrackingTensor):
         return value
+    base = torch.Tensor(value)
     if isinstance(value, pp.LieTensor):
-        unwrapped = torch.Tensor.as_subclass(value, pp.LieTensor)
+        unwrapped = base.as_subclass(pp.LieTensor)
         unwrapped.ltype = value.ltype
         return unwrapped
-    return torch.Tensor.as_subclass(value, torch.Tensor)
+    return base
 
 
 def _unwrap_tracking_tensors(values):
-    if isinstance(values, TrackingTensor):
-        return _unwrap_tracking_tensor(values)
-    if isinstance(values, dict):
-        return {key: _unwrap_tracking_tensors(value) for key, value in values.items()}
-    if isinstance(values, tuple):
-        return tuple(_unwrap_tracking_tensors(value) for value in values)
-    if isinstance(values, list):
-        return [_unwrap_tracking_tensors(value) for value in values]
-    return values
+    return tree_map(_unwrap_tracking_tensor, values)
 
 
 def _rewrap_tracking_tensor(result, tracking_source):
@@ -103,7 +110,7 @@ def _retain_ltype(result, tracking_source, cls, func):
         return result
     if result.shape[-1:] != tracking_source.ltype.dimension:
         return result
-    wrapped = torch.Tensor.as_subclass(result, cls)
+    wrapped = torch.Tensor(result).as_subclass(cls)
     wrapped.ltype = tracking_source.ltype
     return wrapped
 
@@ -178,7 +185,7 @@ def __format__(self, format_spec):
         return format(str(self), format_spec)
 
     def tensor(self) -> torch.Tensor:
-        return torch.Tensor.as_subclass(self, torch.Tensor)
+        return torch.Tensor(self)
 
 
 class _TrackingLieTensor(TrackingTensor, pp.LieTensor):
@@ -195,31 +202,13 @@ def __new__(cls, data, *args, **kwargs):
         return instance
 
     def detach(self):
-        detached = torch.Tensor.as_subclass(super().detach(), type(self))
+        detached = torch.Tensor(self).detach().as_subclass(type(self))
         detached.ltype = self.ltype
         return detached
-"""
-graph design
-Node: (tensor_type: [nn.Parameter, tensor, pp.LieTensor])
-Edge: (indexing, mapping)
-
-G: (V: [Node...], E: [Edge...])
-
-for each e = (u, v) \in E
-parent[loss] = (project, [camera_indexed, point_indexed])
-parent[camera_indexed] = ((indexing, indices), camera_parameters)
-
-build
-parent: key: id(tensor), value: map edge (edge_type, func, [input_args]), index edge (edge_type, indicies, orig_arg)
-
-backward
-1. access loss.parent
-2. check edge type
-3.1. if indexing, permute value
-3.2. if mapping, revise value
-
-recusively call 1-3 until input node is reached. 
-"""
+# Each tracked tensor stores its incoming edge directly as
+# ``(edge_type, edge_metadata, inputs)`` in ``tensor.optrace``. A tensor owns
+# exactly one incoming edge, so an additional dictionary keyed by ``id(tensor)``
+# would be redundant.
 
 
 # =============================================================================