From c2e3762797fca925f1cb0752396ce610082bdc6e Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 00:39:44 +0000 Subject: [PATCH 01/10] dynamo working on froward pass --- README.md | 29 +++++ bae/__init__.py | 2 + bae/autograd/function.py | 40 ++++--- bae/autograd/graph.py | 58 ++++++--- bae/optim/optimizer.py | 37 +++--- bae/utils/pypose_ambient_grad.py | 132 +++++++++++++++++++-- bae/utils/pypose_compile.py | 117 ++++++++++++++++++ tests/autograd/test_bal_jacobian.py | 47 ++++++++ tests/autograd/test_graph_jacobian.py | 4 +- tests/autograd/test_pypose_ambient_grad.py | 13 ++ tests/autograd/test_pypose_compile.py | 100 ++++++++++++++++ 11 files changed, 517 insertions(+), 62 deletions(-) create mode 100644 bae/utils/pypose_compile.py create mode 100644 tests/autograd/test_pypose_compile.py diff --git a/README.md b/README.md index f39f377..9f13cb2 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,35 @@ 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 PyPose + +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. Inductor can therefore load permuted rows directly +inside the fused residual kernels instead of allocating standalone indexed +camera and point tensors. + ## Agent Skills This repo includes skills in [.agent/skills](.agent/skills): diff --git a/bae/__init__.py b/bae/__init__.py index 057e015..9472102 100644 --- a/bae/__init__.py +++ b/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() diff --git a/bae/autograd/function.py b/bae/autograd/function.py index 051df96..43cee3f 100644 --- a/bae/autograd/function.py +++ b/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 @@ -50,15 +51,24 @@ def _attach_cat_trace(result, tensors, dim): if hasattr(t, 'optrace') 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 +79,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): @@ -178,7 +181,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): @@ -210,7 +213,8 @@ def detach(self): 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) +Each tracked tensor stores one ``optrace`` edge directly: a map edge +``(edge_type, func, input_args)`` or index edge ``(edge_type, index, original)``. backward 1. access loss.parent diff --git a/bae/autograd/graph.py b/bae/autograd/graph.py index 6537f87..cf08382 100644 --- a/bae/autograd/graph.py +++ b/bae/autograd/graph.py @@ -8,6 +8,26 @@ from ..sparse import warp_wrappers as _warp_wrappers # noqa: F401 from ..utils.parameter import trim_parameter_jacobian_values +from .function import _INDEXED_TRACE_TAG + + +def _get_optrace(tensor: torch.Tensor): + """Return a trace edge, accepting the pre-Dynamo dictionary format too.""" + trace = tensor.optrace + if isinstance(trace, dict): + return trace[id(tensor)] + return trace + + +def _is_indexed_trace_arg(arg) -> bool: + return isinstance(arg, tuple) and len(arg) == 3 and arg[0] == _INDEXED_TRACE_TAG + + +def _materialize_trace_arg(arg): + if _is_indexed_trace_arg(arg): + _, tensor, index = arg + return tensor[index] + return arg def _crow_to_row_indices(crow_indices: torch.Tensor) -> torch.Tensor: @@ -124,19 +144,24 @@ def _clear_jactrace(output, params): if hasattr(tensor, 'jactrace'): delattr(tensor, 'jactrace') - if not hasattr(tensor, 'optrace') or id(tensor) not in tensor.optrace: + if not hasattr(tensor, 'optrace'): continue - op = tensor.optrace[id(tensor)][0] + trace = _get_optrace(tensor) + op = trace[0] if op == 'map': - args = tensor.optrace[id(tensor)][2] - stack.extend(arg for arg in args if isinstance(arg, torch.Tensor)) + args = trace[2] + for arg in args: + if _is_indexed_trace_arg(arg): + stack.append(arg[1]) + elif isinstance(arg, torch.Tensor): + stack.append(arg) elif op == 'index': - arg = tensor.optrace[id(tensor)][2] + arg = trace[2] if isinstance(arg, torch.Tensor): stack.append(arg) elif op == 'cat': - tracked = tensor.optrace[id(tensor)][2] # ((start, end, arg), ...) + tracked = trace[2] # ((start, end, arg), ...) stack.extend(item[2] for item in tracked if isinstance(item[2], torch.Tensor)) @@ -184,9 +209,10 @@ def backward(output_, is_root=False): if (not is_root) and (not hasattr(output_, 'jactrace')): return - if output_.optrace[id(output_)][0] == 'map': - func = output_.optrace[id(output_)][1] - args = output_.optrace[id(output_)][2] + output_trace = _get_optrace(output_) + if output_trace[0] == 'map': + func = output_trace[1] + args = tuple(_materialize_trace_arg(arg) for arg in output_trace[2]) argnums = tuple(idx for idx, arg in enumerate(args) if hasattr(arg, 'optrace') or isinstance(arg, torch.nn.Parameter)) if len(argnums) == 0: warnings.warn("No upstream parameters to compute jacobian", stacklevel=2) @@ -237,9 +263,9 @@ def backward(output_, is_root=False): delattr(output_, 'jactrace') - elif output_.optrace[id(output_)][0] == 'index': - index = output_.optrace[id(output_)][1] - arg = output_.optrace[id(output_)][2] + elif output_trace[0] == 'index': + index = output_trace[1] + arg = output_trace[2] # If the last operation is indexing, there is no downstream map op to # populate Jacobian values. In this case, the Jacobian block values are @@ -274,9 +300,9 @@ def backward(output_, is_root=False): if hasattr(output_, 'jactrace'): delattr(output_, 'jactrace') - elif output_.optrace[id(output_)][0] == 'cat': - dim = output_.optrace[id(output_)][1] - tracked = output_.optrace[id(output_)][2] # ((start, end, arg), ...) + elif output_trace[0] == 'cat': + dim = output_trace[1] + tracked = output_trace[2] # ((start, end, arg), ...) if dim != 0: raise NotImplementedError("Only torch.cat(..., dim=0) is supported") @@ -314,7 +340,7 @@ def backward(output_, is_root=False): def jacobian(output, params): - assert output.optrace[id(output)][0] in ('map', 'index', 'cat'), "Unsupported last operation in compute graph" + assert _get_optrace(output)[0] in ('map', 'index', 'cat'), "Unsupported last operation in compute graph" _clear_jactrace(output, params) try: backward(output, is_root=True) diff --git a/bae/optim/optimizer.py b/bae/optim/optimizer.py index 19e02b6..f122b4c 100644 --- a/bae/optim/optimizer.py +++ b/bae/optim/optimizer.py @@ -71,20 +71,29 @@ def step(self, input, target=None, weight=None): return self.loss def update_parameter(self, params, step): - numels = [] - 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): - if param.requires_grad: - step_view = d.view(parameter_update_shape(param)) - if getattr(param, 'trim_SE3_grad', False): - param[..., :7] = pp.SE3(param[..., :7]).add_(pp.se3(step_view[..., :6])) - if param.shape[-1] > 7: - param[:, 7:] += step_view[..., 6:] - else: - param.add_(step_view) + with torch.no_grad(): + numels = [] + 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): + if param.requires_grad: + step_view = d.view(parameter_update_shape(param)) + if getattr(param, 'trim_SE3_grad', False): + # 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) + updated_pose = ( + pp.se3(step_view[..., :6]).Exp() + * pp.SE3(param_tensor[..., :7]) + ) + param_tensor[..., :7].copy_(updated_pose.tensor()) + if param.shape[-1] > 7: + param_tensor[..., 7:].add_(step_view[..., 6:]) + else: + param.add_(step_view) class Schur(LM): diff --git a/bae/utils/pypose_ambient_grad.py b/bae/utils/pypose_ambient_grad.py index c51e5b0..759d868 100644 --- a/bae/utils/pypose_ambient_grad.py +++ b/bae/utils/pypose_ambient_grad.py @@ -15,22 +15,48 @@ def _pm(input: torch.Tensor) -> torch.Tensor: return torch.sign(torch.sign(input) * 2 + 1) +def _so3_act_components(qx, qy, qz, qw, px, py, pz): + # Expand p + 2 * qw * (qv x p) + 2 * qv x (qv x p) + # coordinate-wise. This is valid without assuming a unit quaternion and + # avoids reduction/cross operators which otherwise block pointwise fusion. + return ( + (1.0 - 2.0 * qy * qy - 2.0 * qz * qz) * px + + (2.0 * qx * qy - 2.0 * qw * qz) * py + + (2.0 * qx * qz + 2.0 * qw * qy) * pz, + (2.0 * qx * qy + 2.0 * qw * qz) * px + + (1.0 - 2.0 * qx * qx - 2.0 * qz * qz) * py + + (2.0 * qy * qz - 2.0 * qw * qx) * pz, + (2.0 * qx * qz - 2.0 * qw * qy) * px + + (2.0 * qy * qz + 2.0 * qw * qx) * py + + (1.0 - 2.0 * qx * qx - 2.0 * qy * qy) * pz, + ) + + def _so3_act_forward(X: torch.Tensor, p: torch.Tensor) -> torch.Tensor: - Xv, Xw = X[..., :3], X[..., 3:] - uv = torch.linalg.cross(Xv, p, dim=-1) - uv = uv + uv - return p + Xw * uv + torch.linalg.cross(Xv, uv, dim=-1) + quaternion = X.unbind(dim=-1) + point = p.unbind(dim=-1) + return torch.stack(_so3_act_components(*quaternion, *point), dim=-1) def _se3_act_forward(X: torch.Tensor, p: torch.Tensor) -> torch.Tensor: - return X[..., :3] + _so3_act_forward(X[..., 3:], p) + tx, ty, tz, qx, qy, qz, qw = X.unbind(dim=-1) + px, py, pz = p.unbind(dim=-1) + rx, ry, rz = _so3_act_components(qx, qy, qz, qw, px, py, pz) + return torch.stack((tx + rx, ty + ry, tz + rz), dim=-1) def _so3_mul_forward(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor: - Xv, Xw, Yv, Yw = X[..., :3], X[..., 3:], Y[..., :3], Y[..., 3:] - Zv = Xw * Yv + Xv * Yw + torch.linalg.cross(Xv, Yv, dim=-1) - Zw = Xw * Yw - (Xv * Yv).sum(dim=-1, keepdim=True) - return torch.cat([Zv, Zw], dim=-1) + x1, y1, z1, w1 = X.unbind(dim=-1) + x2, y2, z2, w2 = Y.unbind(dim=-1) + return torch.stack( + ( + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ), + dim=-1, + ) def _se3_mul_forward(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor: @@ -67,11 +93,76 @@ def _so3_log_forward(input: torch.Tensor) -> torch.Tensor: return factor * v -def _se3_log_forward(input: torch.Tensor) -> torch.Tensor: - from pypose.lietensor.operation import so3_Jl_inv +def _so3_exp_forward(input: torch.Tensor) -> torch.Tensor: + theta = torch.linalg.norm(input, dim=-1, keepdim=True) + theta2 = theta.square() + theta4 = theta2.square() + nonzero = theta > torch.finfo(theta.dtype).eps + imaginary_factor = nonzero * torch.nan_to_num(torch.sin(0.5 * theta) / theta) + imaginary_factor = imaginary_factor + (~nonzero) * ( + 0.5 - theta2 / 48.0 + theta4 / 3840.0 + ) + real_factor = nonzero * torch.cos(0.5 * theta) + real_factor = real_factor + (~nonzero) * ( + 1.0 - theta2 / 8.0 + theta4 / 384.0 + ) + return torch.cat((input * imaginary_factor, real_factor), dim=-1) + + +def _vec2skew(input: torch.Tensor) -> torch.Tensor: + x, y, z = input.unbind(dim=-1) + zero = torch.zeros_like(x) + return torch.stack( + ( + torch.stack((zero, -z, y), dim=-1), + torch.stack((z, zero, -x), dim=-1), + torch.stack((-y, x, zero), dim=-1), + ), + dim=-2, + ) + + +def _so3_Jl(input: torch.Tensor) -> torch.Tensor: + skew = _vec2skew(input) + theta = torch.linalg.norm(input, dim=-1, keepdim=True).unsqueeze(-1) + theta2 = theta.square() + identity = torch.eye(3, device=input.device, dtype=input.dtype) + identity = identity.expand(input.shape[:-1] + (3, 3)) + nonzero = theta > torch.finfo(theta.dtype).eps + coefficient1 = nonzero * torch.nan_to_num((1.0 - theta.cos()) / theta2) + coefficient1 = coefficient1 + (~nonzero) * (0.5 - theta2 / 24.0) + coefficient2 = nonzero * torch.nan_to_num( + (theta - theta.sin()) / (theta * theta2) + ) + coefficient2 = coefficient2 + (~nonzero) * (1.0 / 6.0 - theta2 / 120.0) + return identity + coefficient1 * skew + coefficient2 * (skew @ skew) + + +def _so3_Jl_inv(input: torch.Tensor) -> torch.Tensor: + skew = _vec2skew(input) + theta = torch.linalg.norm(input, dim=-1, keepdim=True).unsqueeze(-1) + identity = torch.eye(3, device=input.device, dtype=input.dtype) + identity = identity.expand(input.shape[:-1] + (3, 3)) + nonzero = theta > torch.finfo(theta.dtype).eps + coefficient = nonzero * torch.nan_to_num( + (1.0 - theta * (0.5 * theta).cos() / (2.0 * (0.5 * theta).sin())) + / theta.square() + ) + coefficient = coefficient + (~nonzero) * (1.0 / 12.0) + return identity - 0.5 * skew + coefficient * (skew @ skew) + + +def _se3_exp_forward(input: torch.Tensor) -> torch.Tensor: + translation = ( + _so3_Jl(input[..., 3:]) @ input[..., :3].unsqueeze(-1) + ).squeeze(-1) + rotation = _so3_exp_forward(input[..., 3:]) + return torch.cat((translation, rotation), dim=-1) + +def _se3_log_forward(input: torch.Tensor) -> torch.Tensor: phi = _so3_log_forward(input[..., 3:]) - Jl_inv = so3_Jl_inv(phi) + Jl_inv = _so3_Jl_inv(phi) tau = (Jl_inv @ input[..., :3].unsqueeze(-1)).squeeze(-1) return torch.cat([tau, phi], dim=-1) @@ -90,6 +181,13 @@ def install_pypose_ambient_grad_monkeypatch() -> bool: if _PATCH_INSTALLED: return False + from .pypose_compile import install_pypose_torch_compile_monkeypatch + + # The ambient-gradient implementations are plain PyTorch and intended to + # be compiled. Install the semantics-preserving LieTensor compatibility + # shim as part of this patch so the two features remain composable. + install_pypose_torch_compile_monkeypatch() + import pypose.lietensor.lietensor as lt import pypose.lietensor.operation as op @@ -101,6 +199,14 @@ def _ambient_se3_log(self, X): X = X.tensor() if isinstance(X, lt.LieTensor) else X return lt.LieTensor(_se3_log_forward(X), ltype=lt.se3_type) + def _ambient_so3_exp(self, X): + X = X.tensor() if isinstance(X, lt.LieTensor) else X + return lt.LieTensor(_so3_exp_forward(X), ltype=lt.SO3_type) + + def _ambient_se3_exp(self, X): + X = X.tensor() if isinstance(X, lt.LieTensor) else X + return lt.LieTensor(_se3_exp_forward(X), ltype=lt.SE3_type) + def _ambient_so3_act(self, X, p): assert not self.on_manifold and isinstance(p, torch.Tensor) assert p.shape[-1] == 3 or p.shape[-1] == 4, "Invalid Tensor Dimension" @@ -162,6 +268,8 @@ def _ambient_se3_inv(self, X): lt.SO3Type.Log = _ambient_so3_log lt.SE3Type.Log = _ambient_se3_log + lt.so3Type.Exp = _ambient_so3_exp + lt.se3Type.Exp = _ambient_se3_exp lt.SO3Type.Act = _ambient_so3_act lt.SE3Type.Act = _ambient_se3_act lt.SO3Type.Mul = _ambient_so3_mul diff --git a/bae/utils/pypose_compile.py b/bae/utils/pypose_compile.py new file mode 100644 index 0000000..9347cc7 --- /dev/null +++ b/bae/utils/pypose_compile.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import os +import warnings +from importlib import import_module + +import torch +from torch.utils._pytree import tree_flatten, tree_map + + +_ENV_FLAG = "BAE_USE_PYPOSE_TORCH_COMPILE" +_PATCH_INSTALLED = False + + +def pypose_torch_compile_enabled() -> bool: + value = os.environ.get(_ENV_FLAG, "") + return value.lower() in {"1", "true", "yes", "on"} + + +def install_pypose_torch_compile_monkeypatch() -> bool: + """Make PyPose's ``LieTensor`` traceable by TorchDynamo. + + PyPose upcasts a ``LieTensor`` to ``Tensor`` with ``Tensor.as_subclass``. + TorchDynamo cannot trace that direction of ``as_subclass``, while the + equivalent ``Tensor(self)`` alias is supported and keeps autograd history. + PyPose's ``__torch_function__`` also introspects method-wrapper objects with + ``hasattr``, which is another graph break. Its ``vec2skew`` helper creates a + leaf tensor with a dynamic ``requires_grad`` argument, which Dynamo also + rejects; the replacement uses an equivalent constant zero tensor. + + The matmul special case works around Dynamo routing the ``@`` bytecode to + ``Tensor.matmul`` instead of ``LieTensor.__matmul__``. Eager PyPose defines + ``@`` as Lie group multiplication (or action), so preserving that dispatch + is required for compiled code to have the same semantics. + """ + global _PATCH_INSTALLED + if _PATCH_INSTALLED: + return False + + import pypose.lietensor.lietensor as lt + import pypose.lietensor.operation as op + + basics = import_module("pypose.lietensor.basics") + + def _dynamo_compatible_tensor(self): + return torch.Tensor(self) + + def _dynamo_compatible_vec2skew(input): + value = input.tensor() if isinstance(input, lt.LieTensor) else input + assert value.shape[-1] == 3, "Last dim should be 3" + x, y, z = value.unbind(dim=-1) + zero = torch.zeros_like(x) + return torch.stack( + ( + torch.stack((zero, -z, y), dim=-1), + torch.stack((z, zero, -x), dim=-1), + torch.stack((-y, x, zero), dim=-1), + ), + dim=-2, + ) + + @classmethod + def _dynamo_compatible_torch_function(cls, func, types, args=(), kwargs=None): + kwargs = {} if kwargs is None else kwargs + + # Dynamo lowers ``lhs @ rhs`` through Tensor.matmul for a Tensor + # subclass, bypassing LieTensor.__matmul__. + if func is torch.Tensor.matmul: + lhs, rhs = args[:2] + return lhs.ltype.Mul(lhs, rhs) + + tensor_types = tuple( + torch.Tensor + if issubclass(tensor_type, lt.LieTensor) + else tensor_type + for tensor_type in types + ) + data = torch.Tensor.__torch_function__(func, tensor_types, args, kwargs) + + # 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: + flat_args, _ = tree_flatten(args) + ltype = next(arg.ltype for arg in flat_args if isinstance(arg, lt.LieTensor)) + + def wrap(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: + link = "https://pypose.org/docs/main/generated/pypose.LieTensor" + warnings.warn( + f"Tensor Shape Invalid by calling {func}, go to {link}", + stacklevel=2, + ) + return lie_tensor + return tensor + + return tree_map(wrap, data) + return data + + lt.LieTensor.tensor = _dynamo_compatible_tensor + lt.LieTensor.__torch_function__ = _dynamo_compatible_torch_function + # operation.py and lietensor.py import vec2skew directly, so update all + # bound references as well as the defining module. + basics.vec2skew = _dynamo_compatible_vec2skew + op.vec2skew = _dynamo_compatible_vec2skew + lt.vec2skew = _dynamo_compatible_vec2skew + + _PATCH_INSTALLED = True + return True + + +def maybe_install_pypose_torch_compile_monkeypatch() -> bool: + if not pypose_torch_compile_enabled(): + return False + return install_pypose_torch_compile_monkeypatch() diff --git a/tests/autograd/test_bal_jacobian.py b/tests/autograd/test_bal_jacobian.py index 905e8a2..8b73be4 100644 --- a/tests/autograd/test_bal_jacobian.py +++ b/tests/autograd/test_bal_jacobian.py @@ -25,9 +25,14 @@ from pypose.autograd.function import psjac import bae.autograd.graph as autograd_graph # noqa: E402 from bae.optim import LM # noqa: E402 +from bae.utils.pypose_ambient_grad import ( # noqa: E402 + install_pypose_ambient_grad_monkeypatch, +) from bae.utils.pysolvers import PCG # noqa: E402 from datapipes.bal_io import read_bal_data # noqa: E402 +install_pypose_ambient_grad_monkeypatch() + pytestmark = [ pytest.mark.filterwarnings(r"ignore:CUDA initialization.*:UserWarning"), @@ -41,6 +46,48 @@ ("dubrovnik", "problem-356-226730-pre"), ("ladybug", "problem-1723-156502-pre"), ] + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_compiled_indexed_residual_preserves_sparse_jacobian(device: str): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + torch.manual_seed(0) + dtype = torch.float64 + num_cameras, num_points, num_observations = 4, 7, 12 + + cameras = torch.zeros(num_cameras, 10, device=device, dtype=dtype) + cameras[:, 6] = 1.0 + cameras[:, 7] = 500.0 + points = torch.randn(num_points, 3, device=device, dtype=dtype) + points[:, 2].abs_().add_(2.0) + observations = torch.randn(num_observations, 2, device=device, dtype=dtype) + camera_indices = torch.randint(num_cameras, (num_observations,), device=device) + point_indices = torch.randint(num_points, (num_observations,), device=device) + + eager_model = Residual(cameras.clone(), points.clone()).to(device) + expected = eager_model(observations, camera_indices, point_indices) + expected_jacobians = autograd_graph.jacobian( + expected, [eager_model.pose, eager_model.points] + ) + + compiled_model = Residual(cameras.clone(), points.clone()).to(device) + compiled_model = torch.compile(compiled_model, backend="eager", fullgraph=True) + actual = compiled_model(observations, camera_indices, point_indices) + actual_jacobians = autograd_graph.jacobian( + actual, [compiled_model.pose, compiled_model.points] + ) + + torch.testing.assert_close(actual.tensor(), expected.tensor()) + for actual_jacobian, expected_jacobian in zip( + actual_jacobians, expected_jacobians + ): + torch.testing.assert_close( + actual_jacobian.to_dense(), expected_jacobian.to_dense() + ) + + _BAL_TEST_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) _BAL_EXPECTED_FINAL_PER_PIXEL_ERRORS: dict[tuple[str, str], float] = { ("trafalgar", "problem-257-65132-pre"): 0.8588579966685325, diff --git a/tests/autograd/test_graph_jacobian.py b/tests/autograd/test_graph_jacobian.py index ee1f690..9751be3 100644 --- a/tests/autograd/test_graph_jacobian.py +++ b/tests/autograd/test_graph_jacobian.py @@ -424,7 +424,7 @@ def test_pp_parameter_lie_tensor_index_and_cat_preserve_ltype(device: str): assert isinstance(node_a, pp.LieTensor) assert type(node_a.ltype) is type(nodes.ltype) assert hasattr(node_a, "optrace") - assert node_a.optrace[id(node_a)][0] == "index" + assert node_a.optrace[0] == "index" assert node_a.tensor().shape == (idx_a.numel(), 7) assert node_a.translation().shape == (idx_a.numel(), 3) @@ -435,7 +435,7 @@ def test_pp_parameter_lie_tensor_index_and_cat_preserve_ltype(device: str): assert isinstance(cat, pp.LieTensor) assert type(cat.ltype) is type(nodes.ltype) assert hasattr(cat, "optrace") - assert cat.optrace[id(cat)][0] == "cat" + assert cat.optrace[0] == "cat" assert cat.translation().shape == (idx_a.numel() + idx_b.numel(), 3) diff --git a/tests/autograd/test_pypose_ambient_grad.py b/tests/autograd/test_pypose_ambient_grad.py index 8aa7e0f..a50a7d0 100644 --- a/tests/autograd/test_pypose_ambient_grad.py +++ b/tests/autograd/test_pypose_ambient_grad.py @@ -83,3 +83,16 @@ def test_pypose_se3_inv_monkeypatch_produces_ambient_gradient(): jac_fd = _fd_jacobian(lambda x: pp.SE3(x).Inv().tensor(), pose) torch.testing.assert_close(jac, jac_fd, rtol=1e-7, atol=1e-7) + + +def test_pypose_se3_exp_monkeypatch_produces_ambient_gradient(): + install_pypose_ambient_grad_monkeypatch() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float64 + + algebra = 0.2 * torch.randn(1, 6, device=device, dtype=dtype) + + jac = torch.func.jacrev(lambda x: pp.se3(x).Exp().tensor())(algebra) + jac_fd = _fd_jacobian(lambda x: pp.se3(x).Exp().tensor(), algebra) + + torch.testing.assert_close(jac, jac_fd, rtol=1e-7, atol=1e-7) diff --git a/tests/autograd/test_pypose_compile.py b/tests/autograd/test_pypose_compile.py new file mode 100644 index 0000000..8499a56 --- /dev/null +++ b/tests/autograd/test_pypose_compile.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import pypose as pp +import torch + +from bae.utils.pypose_ambient_grad import install_pypose_ambient_grad_monkeypatch +from bae.utils.pypose_compile import install_pypose_torch_compile_monkeypatch +from bae.optim.optimizer import LM + + +def test_lietensor_tensor_alias_is_fullgraph_traceable_and_differentiable(): + install_pypose_torch_compile_monkeypatch() + data = pp.randn_SE3(4, sigma=0.2, dtype=torch.float64).tensor().detach().requires_grad_() + pose = pp.SE3(data) + + compiled = torch.compile( + lambda value: value.tensor().square().sum(), backend="eager", fullgraph=True + ) + loss = compiled(pose) + (gradient,) = torch.autograd.grad(loss, data) + + torch.testing.assert_close(gradient, 2.0 * data) + + compiled_inverse = torch.compile( + lambda value: value.Inv(), backend="eager", fullgraph=True + ) + inverse = compiled_inverse(pose) + assert isinstance(inverse, pp.LieTensor) + assert inverse.ltype is pp.SE3_type + + +def test_ambient_se3_pipeline_is_fullgraph_traceable_with_local_gradients(): + install_pypose_ambient_grad_monkeypatch() + dtype = torch.float64 + lhs = pp.randn_SE3(8, sigma=0.2, dtype=dtype).tensor().detach().requires_grad_() + rhs = pp.randn_SE3(8, sigma=0.2, dtype=dtype).tensor().detach().requires_grad_() + points = torch.randn(8, 3, dtype=dtype, requires_grad=True) + + def residual(a, b, p): + pose_a = pp.SE3(a) + pose_b = pp.SE3(b) + relative = (pose_a.Inv() @ pose_b).Log().tensor() + return torch.cat((relative, pose_a.Act(p)), dim=-1) + + eager = residual(lhs, rhs, points) + eager_gradients = torch.autograd.grad( + eager.square().sum(), (lhs, rhs, points), retain_graph=True + ) + + compiled = torch.compile(residual, backend="eager", fullgraph=True) + actual = compiled(lhs, rhs, points) + compiled_gradients = torch.autograd.grad(actual.square().sum(), (lhs, rhs, points)) + + torch.testing.assert_close(actual, eager) + for actual_gradient, eager_gradient in zip(compiled_gradients, eager_gradients): + torch.testing.assert_close(actual_gradient, eager_gradient) + + +def test_ambient_se3_jacrev_is_fullgraph_traceable(): + install_pypose_ambient_grad_monkeypatch() + pose = pp.randn_SE3(3, sigma=0.2, dtype=torch.float64).tensor() + points = torch.randn(3, 3, dtype=torch.float64) + + jacobian = torch.func.jacrev(lambda value: pp.SE3(value).Act(points)) + expected = jacobian(pose) + compiled_jacobian = torch.compile(jacobian, backend="eager", fullgraph=True) + + torch.testing.assert_close(compiled_jacobian(pose), expected) + + algebra = 0.2 * torch.randn(3, 6, dtype=torch.float64) + exp_jacobian = torch.func.jacrev(lambda value: pp.se3(value).Exp().tensor()) + expected_exp = exp_jacobian(algebra) + compiled_exp = torch.compile(exp_jacobian, backend="eager", fullgraph=True) + + torch.testing.assert_close(compiled_exp(algebra), expected_exp) + + +def test_lm_trimmed_se3_update_composes_with_compile_patch(): + install_pypose_ambient_grad_monkeypatch() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float64 + + data = torch.zeros(3, 10, device=device, dtype=dtype) + data[:, 6] = 1.0 + data[:, 7:] = torch.randn(3, 3, device=device, dtype=dtype) + parameter = pp.Parameter(data.clone(), sjac=True) + parameter.trim_SE3_grad = True + update = 0.01 * torch.randn(3, 9, device=device, dtype=dtype) + + expected_pose = ( + pp.se3(update[:, :6]).Exp() * pp.SE3(data[:, :7]) + ).tensor() + expected_intrinsics = data[:, 7:] + update[:, 6:] + + # update_parameter does not inspect optimizer state, so exercise it as an + # unbound method without constructing a linear solver. + LM.update_parameter(None, [parameter], update.flatten()) + + torch.testing.assert_close(parameter.tensor()[:, :7], expected_pose) + torch.testing.assert_close(parameter.tensor()[:, 7:], expected_intrinsics) From aad651ddcff41abbfb97a5e8a0fc6f7218140232 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 00:44:34 +0000 Subject: [PATCH 02/10] backward works in dynamo --- README.md | 37 +++- bae/autograd/function.py | 12 +- bae/autograd/graph.py | 268 +++++++++++++++++++++++++- bae/utils/parameter.py | 4 +- tests/autograd/test_bal_jacobian.py | 23 ++- tests/autograd/test_graph_jacobian.py | 39 ++++ 6 files changed, 361 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 9f13cb2..977855d 100644 --- a/README.md +++ b/README.md @@ -158,9 +158,40 @@ 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. Inductor can therefore load permuted rows directly -inside the fused residual kernels instead of allocating standalone indexed -camera and point tensors. +to escape the compiled graph. This gives Inductor the opportunity to load +permuted rows directly inside fused kernels, but does not guarantee that it will +do so. Inductor may instead materialize standalone indexed tensors when the +gathered values have multiple downstream consumers, as can happen during +Jacobian computation. `fullgraph=True` guarantees graph capture rather than a +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 diff --git a/bae/autograd/function.py b/bae/autograd/function.py index 43cee3f..4b2946c 100644 --- a/bae/autograd/function.py +++ b/bae/autograd/function.py @@ -48,7 +48,11 @@ 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 = ("cat", dim, tuple(tracked)) @@ -106,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 @@ -198,10 +202,10 @@ 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 -""" +r""" graph design Node: (tensor_type: [nn.Parameter, tensor, pp.LieTensor]) Edge: (indexing, mapping) diff --git a/bae/autograd/graph.py b/bae/autograd/graph.py index cf08382..b9e4a62 100644 --- a/bae/autograd/graph.py +++ b/bae/autograd/graph.py @@ -1,5 +1,5 @@ -from typing import Optional +from typing import NamedTuple, Optional import warnings import pypose as pp @@ -11,6 +11,21 @@ from .function import _INDEXED_TRACE_TAG +class BSRJacobianData(NamedTuple): + """Dense tensor components of a BSR Jacobian. + + PyTorch's FakeTensor/AOTAutograd stack cannot currently represent a sparse + BSR graph output. Returning its component tensors keeps Jacobian traversal + and value computation inside the compiled graph; materialization into a + sparse Tensor is a zero-copy eager wrapper operation after the call. + """ + + crow_indices: torch.Tensor + col_indices: torch.Tensor + values: torch.Tensor + size: tuple[int, int] + + def _get_optrace(tensor: torch.Tensor): """Return a trace edge, accepting the pre-Dynamo dictionary format too.""" trace = tensor.optrace @@ -24,6 +39,7 @@ def _is_indexed_trace_arg(arg) -> bool: def _materialize_trace_arg(arg): + # produces an ordinary aten.index; it enables Inductor to choose fusion but does not force it. if _is_indexed_trace_arg(arg): _, tensor, index = arg return tensor[index] @@ -37,12 +53,12 @@ def _crow_to_row_indices(crow_indices: torch.Tensor) -> torch.Tensor: def _row_indices_to_crow(row_indices: torch.Tensor, n_rows: int, dtype: torch.dtype) -> torch.Tensor: - if row_indices.numel() == 0: - return torch.zeros(n_rows + 1, device=row_indices.device, dtype=dtype) - counts = torch.bincount(row_indices.to(torch.int64), minlength=n_rows).to(dtype) - crow = torch.zeros(n_rows + 1, device=row_indices.device, dtype=dtype) - crow[1:] = torch.cumsum(counts, dim=0) - return crow + row_boundaries = torch.arange( + n_rows + 1, + device=row_indices.device, + dtype=row_indices.dtype, + ) + return torch.searchsorted(row_indices, row_boundaries).to(dtype) def _slice_upstream_bsr_columns( @@ -129,8 +145,7 @@ def construct_sbt(jac_from_vmap, num, index: Optional[torch.Tensor], type=torch. return torch.sparse_bsr_tensor(col_indices=index, crow_indices=torch.arange(n + 1, device=index.device, dtype=idx_dtype), values = jac_from_vmap, - size = (n * block_shape[0], num * block_shape[1]), - device=index.device, dtype=jac_from_vmap.dtype) + size = (n * block_shape[0], num * block_shape[1])) def _clear_jactrace(output, params): seen = set() @@ -203,6 +218,241 @@ def _vmap_in_dims(args): """Map Tensor inputs over dim 0 while treating non-Tensors as constants.""" return tuple(0 if isinstance(arg, torch.Tensor) else None for arg in args) + +def _restore_lie_inputs(func, args): + """Reconstruct LieTensor inputs after functorch strips subclass metadata.""" + ltypes = tuple( + arg.ltype if isinstance(arg, pp.LieTensor) else None for arg in args + ) + if not any(ltype is not None for ltype in ltypes): + return func + + def wrapped(*values): + values = tuple( + pp.LieTensor(value, ltype=ltype) if ltype is not None else value + for value, ltype in zip(values, ltypes) + ) + return func(*values) + + return wrapped + + +def _compose_map_trace(jac_block, upstream): + """Compose a local block Jacobian with an optional downstream trace.""" + if upstream is None: + rows = torch.arange( + jac_block.shape[0], device=jac_block.device, dtype=torch.int64 + ) + return (rows, None, jac_block) + + rows, positions, upstream_values = upstream + if positions is not None: + jac_block = jac_block[positions] + return (rows, positions, upstream_values @ jac_block) + + +def _compose_index_trace(index, upstream, output): + """Propagate a block trace through an observation-wise index operation.""" + if upstream is None: + rows = torch.arange( + output.shape[0], device=output.device, dtype=torch.int64 + ) + if output.ndim == 1: + values = torch.ones( + (output.shape[0], 1, 1), + device=output.device, + dtype=output.dtype, + ) + else: + 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) + return (rows, index, values) + + rows, positions, values = upstream + if positions is not None: + index = index[positions] + return (rows, index, values) + + +def _parameter_index(tensor, params): + """Resolve a traced leaf to its requested parameter using object identity.""" + for index, param in enumerate(params): + if tensor is param: + return index + return None + + +def _append_functional_trace(tensor, trace, params, contributions): + """Continue through a traced node or record a parameter contribution.""" + if isinstance(tensor, torch.Tensor) and hasattr(tensor, "optrace"): + _functional_backward(tensor, trace, params, contributions) + return + + index = _parameter_index(tensor, params) + if index is not None: + contributions[index].append(trace) + + +def _functional_backward(output, upstream, params, contributions): + """Compile-safe sparse-Jacobian traversal without Tensor attribute mutation. + + Dynamo unrolls the Python control flow over the immutable ``optrace`` tree. + Runtime Jacobian state is represented only by tensor tuples, rather than by + adding and deleting ``jactrace`` attributes on Tensor/FakeTensor objects. + """ + output_trace = _get_optrace(output) + op = output_trace[0] + + if op == "map": + func = output_trace[1] + trace_args = output_trace[2] + # This design avoids retaining the intermediate indexed TrackingTensor in the immutable trace tree. Under torch.compile, the + # reconstructed cameras[camera_indices] is captured as an ordinary graph operation, so Inductor may fuse its indexed loads with + # the derivative computation. + args = tuple(_materialize_trace_arg(arg) for arg in trace_args) + argnums = tuple( + index + for index, arg in enumerate(args) + if hasattr(arg, "optrace") or isinstance(arg, torch.nn.Parameter) + ) + if len(argnums) == 0: + return + + jac_blocks = torch.vmap( + jacrev(_restore_lie_inputs(func, args), argnums=argnums), + in_dims=_vmap_in_dims(args), + )(*args) + + for jac_block, arg_index in zip(jac_blocks, argnums): + assert jac_block.ndim == 3, "`func` is not properly vectorized in `torch.vmap`" + trace = _compose_map_trace(jac_block, upstream) + trace_arg = trace_args[arg_index] + if _is_indexed_trace_arg(trace_arg): + _, parent, index = trace_arg + trace = _compose_index_trace(index, trace, args[arg_index]) + _append_functional_trace(parent, trace, params, contributions) + else: + _append_functional_trace( + args[arg_index], trace, params, contributions + ) + return + + if op == "index": + index, parent = output_trace[1], output_trace[2] + trace = _compose_index_trace(index, upstream, output) + _append_functional_trace(parent, trace, params, contributions) + return + + if op == "cat": + dim, tracked = output_trace[1], output_trace[2] + if dim != 0: + raise NotImplementedError("Only torch.cat(..., dim=0) is supported") + + if upstream is None: + rows = torch.arange( + output.shape[0], device=output.device, dtype=torch.int64 + ) + positions = rows + if output.ndim == 1: + values = torch.ones( + (output.shape[0], 1, 1), + device=output.device, + dtype=output.dtype, + ) + else: + 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) + else: + rows, positions, values = upstream + if positions is None: + positions = torch.arange( + values.shape[0], device=values.device, dtype=torch.int64 + ) + + for start, end, arg in tracked: + mask = (positions >= start) & (positions < end) + trace = ( + rows[mask], + positions[mask] - start, + values[mask], + ) + _append_functional_trace(arg, trace, params, contributions) + return + + raise NotImplementedError(f"Unsupported trace operation: {op}") + + +def jacobian_components(output, params): + """Return compile-safe BSR component tensors for each requested parameter. + + The result contains one or more components per parameter. Multiple + components arise when distinct compute-graph branches contribute to the + same parameter and are combined during eager materialization. + """ + assert _get_optrace(output)[0] in ( + "map", + "index", + "cat", + ), "Unsupported last operation in compute graph" + contributions = [[] for _ in params] + _functional_backward(output, None, params, contributions) + + result = [] + for param, traces in zip(params, contributions): + param_components = [] + for rows, columns, values in traces: + if columns is None: + columns = torch.arange( + values.shape[0], device=values.device, dtype=torch.int64 + ) + values = trim_parameter_jacobian_values( + param, values, block_indices=columns + ) + crow = _row_indices_to_crow( + rows, + output.shape[0], + dtype=columns.dtype, + ) + param_components.append( + BSRJacobianData( + crow, + columns, + values, + ( + output.shape[0] * values.shape[-2], + param.shape[0] * values.shape[-1], + ), + ) + ) + result.append(tuple(param_components)) + return tuple(result) + + +def materialize_jacobian_components(components): + """Create sparse BSR tensors from ``jacobian_components`` output.""" + result = [] + for param_components in components: + if len(param_components) == 0: + continue + sparse_components = [ + torch.sparse_bsr_tensor( + component.crow_indices, + component.col_indices, + component.values, + size=component.size, + ) + for component in param_components + ] + combined = sparse_components[0] + for sparse_component in sparse_components[1:]: + combined = combined + sparse_component + result.append(combined) + return result + def backward(output_, is_root=False): # For non-root recursion, no incoming trace means no contribution to # propagate. This avoids re-initializing identity traces on revisits. diff --git a/bae/utils/parameter.py b/bae/utils/parameter.py index 42bb1de..580d513 100644 --- a/bae/utils/parameter.py +++ b/bae/utils/parameter.py @@ -25,7 +25,7 @@ def trim_parameter_jacobian_values( if getattr(param, 'trim_SE3_grad', False): 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() if block_indices is not None: pose = pose[block_indices.to(torch.long)] pose_values = values[..., :7] @ se3_retraction_jacobian(pose) @@ -34,7 +34,7 @@ def trim_parameter_jacobian_values( return torch.cat([pose_values, values[..., 7:]], dim=-1) if isinstance(param, pp.LieTensor): if pypose_ambient_grad_enabled(): - lie_param = param.detach() + lie_param = torch.Tensor(param).detach() if block_indices is not None: lie_param = lie_param[block_indices.to(torch.long)] if param.ltype == pp.SO3_type: diff --git a/tests/autograd/test_bal_jacobian.py b/tests/autograd/test_bal_jacobian.py index 8b73be4..34dd9d1 100644 --- a/tests/autograd/test_bal_jacobian.py +++ b/tests/autograd/test_bal_jacobian.py @@ -73,11 +73,26 @@ def test_compiled_indexed_residual_preserves_sparse_jacobian(device: str): ) compiled_model = Residual(cameras.clone(), points.clone()).to(device) - compiled_model = torch.compile(compiled_model, backend="eager", fullgraph=True) - actual = compiled_model(observations, camera_indices, point_indices) - actual_jacobians = autograd_graph.jacobian( - actual, [compiled_model.pose, compiled_model.points] + + def residual_and_jacobian_components(observations, camera_indices, point_indices): + residual = compiled_model(observations, camera_indices, point_indices) + components = autograd_graph.jacobian_components( + residual, (compiled_model.pose, compiled_model.points) + ) + return residual, components + + backend = "inductor" if device == "cuda" else "eager" + compiled = torch.compile( + residual_and_jacobian_components, + backend=backend, + fullgraph=True, ) + actual, components = compiled(observations, camera_indices, point_indices) + actual_jacobians = autograd_graph.materialize_jacobian_components(components) + + for component_group, actual_jacobian in zip(components, actual_jacobians): + assert len(component_group) == 1 + assert component_group[0].values.data_ptr() == actual_jacobian.values().data_ptr() torch.testing.assert_close(actual.tensor(), expected.tensor()) for actual_jacobian, expected_jacobian in zip( diff --git a/tests/autograd/test_graph_jacobian.py b/tests/autograd/test_graph_jacobian.py index 9751be3..959fa89 100644 --- a/tests/autograd/test_graph_jacobian.py +++ b/tests/autograd/test_graph_jacobian.py @@ -10,6 +10,7 @@ from pypose.autograd.function import psjac from bae.autograd.graph import jacobian as sparse_jacobian +from bae.autograd.graph import jacobian_components, materialize_jacobian_components from bae.utils.retraction_jacobian import se3_retraction_jacobian from bae.utils.pypose_ambient_grad import ( install_pypose_ambient_grad_monkeypatch, @@ -296,6 +297,17 @@ def f(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: assert JB_sparse.crow_indices()[-1].item() == n_b assert torch.equal(JB_sparse.col_indices(), idx_b) + def compiled_components(*args): + residual = model(*args) + return jacobian_components(residual, (model.A, model.B)) + + components = torch.compile( + compiled_components, backend="eager", fullgraph=True + )(obs_a, obs_b, idx_a, idx_b, mul_a, mul_b) + compiled_jacobians = materialize_jacobian_components(components) + torch.testing.assert_close(compiled_jacobians[0].to_dense(), JA_sparse.to_dense()) + torch.testing.assert_close(compiled_jacobians[1].to_dense(), JB_sparse.to_dense()) + class CatSubResidual(nn.Module): def __init__(self, A: torch.Tensor, B: torch.Tensor): @@ -403,6 +415,17 @@ def f(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: torch.testing.assert_close(JA_sparse.to_dense(), _flatten_jac(JA), rtol=1e-10, atol=1e-10) torch.testing.assert_close(JB_sparse.to_dense(), _flatten_jac(JB), rtol=1e-10, atol=1e-10) + def compiled_components(obs, index): + residual = model(obs, index) + return jacobian_components(residual, (model.A, model.B)) + + components = torch.compile( + compiled_components, backend="eager", fullgraph=True + )(obs, idx) + compiled_jacobians = materialize_jacobian_components(components) + torch.testing.assert_close(compiled_jacobians[0].to_dense(), JA_sparse.to_dense()) + torch.testing.assert_close(compiled_jacobians[1].to_dense(), JB_sparse.to_dense()) + @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_pp_parameter_lie_tensor_index_and_cat_preserve_ltype(device: str): @@ -471,3 +494,19 @@ def f(nodes_data: torch.Tensor) -> torch.Tensor: else: J_ref = _flatten_jac(J_dense[..., :6]) torch.testing.assert_close(J_sparse.to_dense(), J_ref, rtol=1e-10, atol=1e-10) + + def compiled_components(poses, first_indices, second_indices): + residual = _relative_se3_residual( + poses, + model[first_indices], + model[second_indices], + ) + return jacobian_components(residual, (model,)) + + components = torch.compile( + compiled_components, backend="eager", fullgraph=True + )(poses, idx1, idx2) + (compiled_jacobian,) = materialize_jacobian_components(components) + torch.testing.assert_close( + compiled_jacobian.to_dense(), J_sparse.to_dense(), rtol=1e-10, atol=1e-10 + ) From 17ab95724ded54b2abb410c9a5c5331a4cdc44cb Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 00:50:09 +0000 Subject: [PATCH 03/10] update plan --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 977855d..c1effa8 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,7 @@ ### 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)) From 300d64cb61d6661e2f38553b2862f40c89bd94e2 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 00:50:51 +0000 Subject: [PATCH 04/10] reorder --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c1effa8..80794be 100644 --- a/README.md +++ b/README.md @@ -82,11 +82,11 @@ - **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 (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 From bdd3210d521bfd65f592d264aff4d74e9efa3ef5 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 00:51:36 +0000 Subject: [PATCH 05/10] remove AMGX plan as distributed PCG is finished --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index 80794be..11d2d7f 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,6 @@ ### Future Plan - [ ] 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) From 64aad04d50da667e46fdc261def1499b1d2cbff9 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 04:27:42 +0000 Subject: [PATCH 06/10] =?UTF-8?q?deprecate=20the=20dict=20type=20optrace.?= =?UTF-8?q?=20-=20The=20dictionary=20form=20was=20unnecessary.=20Each=20te?= =?UTF-8?q?nsor=20owns=20exactly=20one=20incoming=20edge,=20so=20id(tensor?= =?UTF-8?q?)=20=E2=86=92=20edge=20was=20redundant.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bae/autograd/function.py | 27 ++++----------------------- bae/autograd/graph.py | 18 +++++------------- tests/autograd/test_graph_jacobian.py | 4 ++++ 3 files changed, 13 insertions(+), 36 deletions(-) diff --git a/bae/autograd/function.py b/bae/autograd/function.py index 4b2946c..da506e7 100644 --- a/bae/autograd/function.py +++ b/bae/autograd/function.py @@ -205,29 +205,10 @@ def detach(self): detached = torch.Tensor(self).detach().as_subclass(type(self)) detached.ltype = self.ltype return detached -r""" -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 -Each tracked tensor stores one ``optrace`` edge directly: a map edge -``(edge_type, func, input_args)`` or index edge ``(edge_type, index, original)``. - -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. # ============================================================================= diff --git a/bae/autograd/graph.py b/bae/autograd/graph.py index b9e4a62..29e2915 100644 --- a/bae/autograd/graph.py +++ b/bae/autograd/graph.py @@ -26,14 +26,6 @@ class BSRJacobianData(NamedTuple): size: tuple[int, int] -def _get_optrace(tensor: torch.Tensor): - """Return a trace edge, accepting the pre-Dynamo dictionary format too.""" - trace = tensor.optrace - if isinstance(trace, dict): - return trace[id(tensor)] - return trace - - def _is_indexed_trace_arg(arg) -> bool: return isinstance(arg, tuple) and len(arg) == 3 and arg[0] == _INDEXED_TRACE_TAG @@ -162,7 +154,7 @@ def _clear_jactrace(output, params): if not hasattr(tensor, 'optrace'): continue - trace = _get_optrace(tensor) + trace = tensor.optrace op = trace[0] if op == 'map': args = trace[2] @@ -301,7 +293,7 @@ def _functional_backward(output, upstream, params, contributions): Runtime Jacobian state is represented only by tensor tuples, rather than by adding and deleting ``jactrace`` attributes on Tensor/FakeTensor objects. """ - output_trace = _get_optrace(output) + output_trace = output.optrace op = output_trace[0] if op == "map": @@ -393,7 +385,7 @@ def jacobian_components(output, params): components arise when distinct compute-graph branches contribute to the same parameter and are combined during eager materialization. """ - assert _get_optrace(output)[0] in ( + assert output.optrace[0] in ( "map", "index", "cat", @@ -459,7 +451,7 @@ def backward(output_, is_root=False): if (not is_root) and (not hasattr(output_, 'jactrace')): return - output_trace = _get_optrace(output_) + output_trace = output_.optrace if output_trace[0] == 'map': func = output_trace[1] args = tuple(_materialize_trace_arg(arg) for arg in output_trace[2]) @@ -590,7 +582,7 @@ def backward(output_, is_root=False): def jacobian(output, params): - assert _get_optrace(output)[0] in ('map', 'index', 'cat'), "Unsupported last operation in compute graph" + assert output.optrace[0] in ('map', 'index', 'cat'), "Unsupported last operation in compute graph" _clear_jactrace(output, params) try: backward(output, is_root=True) diff --git a/tests/autograd/test_graph_jacobian.py b/tests/autograd/test_graph_jacobian.py index 959fa89..c4de2d2 100644 --- a/tests/autograd/test_graph_jacobian.py +++ b/tests/autograd/test_graph_jacobian.py @@ -94,6 +94,8 @@ def test_sparse_jacobian_matches_torch_jacrev(device: str): model = ToyResidual(A0, B0) out = model(obs, idx_a, idx_b, sel) + assert isinstance(out.optrace, tuple) + assert out.optrace[0] == "map" J_sparse = sparse_jacobian(out, [model.A, model.B]) assert len(J_sparse) == 2 @@ -447,6 +449,7 @@ def test_pp_parameter_lie_tensor_index_and_cat_preserve_ltype(device: str): assert isinstance(node_a, pp.LieTensor) assert type(node_a.ltype) is type(nodes.ltype) assert hasattr(node_a, "optrace") + assert isinstance(node_a.optrace, tuple) assert node_a.optrace[0] == "index" assert node_a.tensor().shape == (idx_a.numel(), 7) @@ -458,6 +461,7 @@ def test_pp_parameter_lie_tensor_index_and_cat_preserve_ltype(device: str): assert isinstance(cat, pp.LieTensor) assert type(cat.ltype) is type(nodes.ltype) assert hasattr(cat, "optrace") + assert isinstance(cat.optrace, tuple) assert cat.optrace[0] == "cat" assert cat.translation().shape == (idx_a.numel() + idx_b.numel(), 3) From 3eb722cc682460b16ea7a46830fe371d85b2386f Mon Sep 17 00:00:00 2001 From: Zihang Fang <134358308+ZihangFang@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:44:06 -0700 Subject: [PATCH 07/10] Update README.md --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 11d2d7f..86de965 100644 --- a/README.md +++ b/README.md @@ -132,7 +132,7 @@ 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 PyPose +### `torch.compile` with `LieTensor` Enable the `LieTensor` TorchDynamo compatibility shim before importing either `pypose` or `bae`: @@ -158,10 +158,10 @@ 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, but does not guarantee that it will -do so. Inductor may instead materialize standalone indexed tensors when the +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 rather than a +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 From 54081a117530d18472cbd0f0039cb4161954c3b3 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 19:55:32 +0000 Subject: [PATCH 08/10] fix the none comment --- bae/autograd/graph.py | 7 +++++- tests/autograd/test_graph_jacobian.py | 35 +++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/bae/autograd/graph.py b/bae/autograd/graph.py index 29e2915..f861faa 100644 --- a/bae/autograd/graph.py +++ b/bae/autograd/graph.py @@ -425,10 +425,15 @@ def jacobian_components(output, params): def materialize_jacobian_components(components): - """Create sparse BSR tensors from ``jacobian_components`` output.""" + """Create sparse BSR tensors while preserving parameter alignment. + + An empty component group represents a parameter that does not contribute + to the output and is materialized as ``None``. + """ result = [] for param_components in components: if len(param_components) == 0: + result.append(None) continue sparse_components = [ torch.sparse_bsr_tensor( diff --git a/tests/autograd/test_graph_jacobian.py b/tests/autograd/test_graph_jacobian.py index c4de2d2..0e63cc7 100644 --- a/tests/autograd/test_graph_jacobian.py +++ b/tests/autograd/test_graph_jacobian.py @@ -116,6 +116,41 @@ def f(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: assert torch.equal(J_sparse[1].col_indices(), idx_b[sel]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_materialized_components_preserve_unused_parameter_position(device: str): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + torch.manual_seed(0) + dtype = torch.float64 + dim = 3 + + model = ToyResidual( + torch.randn(4, dim, device=device, dtype=dtype), + torch.randn(5, dim, device=device, dtype=dtype), + ) + unused = pp.Parameter(torch.randn(6, dim, device=device, dtype=dtype), sjac=True) + obs = torch.randn(7, dim, device=device, dtype=dtype) + idx_a = torch.randint(4, (7,), device=device, dtype=torch.int32) + idx_b = torch.randint(5, (7,), device=device, dtype=torch.int32) + sel = torch.tensor([0, 2, 5, 6], device=device, dtype=torch.int32) + + def compiled_components(obs, idx_a, idx_b, sel): + residual = model(obs, idx_a, idx_b, sel) + return jacobian_components(residual, (model.A, unused, model.B)) + + components = torch.compile( + compiled_components, backend="eager", fullgraph=True + )(obs, idx_a, idx_b, sel) + jacobians = materialize_jacobian_components(components) + + assert len(components) == 3 + assert len(jacobians) == 3 + assert jacobians[0] is not None + assert jacobians[1] is None + assert jacobians[2] is not None + + @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_sparse_jacobian_last_op_indexing_is_identity(device: str): if device == "cuda" and not torch.cuda.is_available(): From e2ce06830914dd5dd2e1f025dcdd2c547c862143 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 20:00:42 +0000 Subject: [PATCH 09/10] fix torch.where comments --- bae/utils/pypose_ambient_grad.py | 45 ++++++++++++++++++--------- tests/autograd/test_pypose_compile.py | 39 ++++++++++++++++++++++- 2 files changed, 68 insertions(+), 16 deletions(-) diff --git a/bae/utils/pypose_ambient_grad.py b/bae/utils/pypose_ambient_grad.py index 759d868..1ff3642 100644 --- a/bae/utils/pypose_ambient_grad.py +++ b/bae/utils/pypose_ambient_grad.py @@ -98,13 +98,16 @@ def _so3_exp_forward(input: torch.Tensor) -> torch.Tensor: theta2 = theta.square() theta4 = theta2.square() nonzero = theta > torch.finfo(theta.dtype).eps - imaginary_factor = nonzero * torch.nan_to_num(torch.sin(0.5 * theta) / theta) - imaginary_factor = imaginary_factor + (~nonzero) * ( - 0.5 - theta2 / 48.0 + theta4 / 3840.0 + 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 = nonzero * torch.cos(0.5 * theta) - real_factor = real_factor + (~nonzero) * ( - 1.0 - theta2 / 8.0 + theta4 / 384.0 + real_factor = torch.where( + nonzero, + torch.cos(0.5 * theta), + 1.0 - theta2 / 8.0 + theta4 / 384.0, ) return torch.cat((input * imaginary_factor, real_factor), dim=-1) @@ -129,12 +132,18 @@ def _so3_Jl(input: torch.Tensor) -> torch.Tensor: identity = torch.eye(3, device=input.device, dtype=input.dtype) identity = identity.expand(input.shape[:-1] + (3, 3)) nonzero = theta > torch.finfo(theta.dtype).eps - coefficient1 = nonzero * torch.nan_to_num((1.0 - theta.cos()) / theta2) - coefficient1 = coefficient1 + (~nonzero) * (0.5 - theta2 / 24.0) - coefficient2 = nonzero * torch.nan_to_num( - (theta - theta.sin()) / (theta * theta2) + safe_theta = torch.where(nonzero, theta, torch.ones_like(theta)) + safe_theta2 = safe_theta.square() + coefficient1 = torch.where( + nonzero, + (1.0 - safe_theta.cos()) / safe_theta2, + 0.5 - theta2 / 24.0, + ) + coefficient2 = torch.where( + nonzero, + (safe_theta - safe_theta.sin()) / (safe_theta * safe_theta2), + 1.0 / 6.0 - theta2 / 120.0, ) - coefficient2 = coefficient2 + (~nonzero) * (1.0 / 6.0 - theta2 / 120.0) return identity + coefficient1 * skew + coefficient2 * (skew @ skew) @@ -144,11 +153,17 @@ def _so3_Jl_inv(input: torch.Tensor) -> torch.Tensor: identity = torch.eye(3, device=input.device, dtype=input.dtype) identity = identity.expand(input.shape[:-1] + (3, 3)) nonzero = theta > torch.finfo(theta.dtype).eps - coefficient = nonzero * torch.nan_to_num( - (1.0 - theta * (0.5 * theta).cos() / (2.0 * (0.5 * theta).sin())) - / theta.square() + safe_theta = torch.where(nonzero, theta, torch.ones_like(theta)) + half_safe_theta = 0.5 * safe_theta + coefficient = torch.where( + nonzero, + ( + 1.0 + - half_safe_theta * half_safe_theta.cos() / half_safe_theta.sin() + ) + / safe_theta.square(), + torch.full_like(theta, 1.0 / 12.0), ) - coefficient = coefficient + (~nonzero) * (1.0 / 12.0) return identity - 0.5 * skew + coefficient * (skew @ skew) diff --git a/tests/autograd/test_pypose_compile.py b/tests/autograd/test_pypose_compile.py index 8499a56..e5ee3e2 100644 --- a/tests/autograd/test_pypose_compile.py +++ b/tests/autograd/test_pypose_compile.py @@ -1,13 +1,50 @@ from __future__ import annotations import pypose as pp +import pytest import torch -from bae.utils.pypose_ambient_grad import install_pypose_ambient_grad_monkeypatch +from bae.utils.pypose_ambient_grad import ( + _so3_exp_forward, + _so3_Jl, + _so3_Jl_inv, + install_pypose_ambient_grad_monkeypatch, +) from bae.utils.pypose_compile import install_pypose_torch_compile_monkeypatch from bae.optim.optimizer import LM +@torch.no_grad() +def _zero_and_small_angles(dtype: torch.dtype) -> torch.Tensor: + inputs = torch.zeros(2, 3, dtype=dtype) + inputs[1, 0] = torch.finfo(dtype).eps / 2.0 + return inputs + + +@pytest.mark.parametrize("function", (_so3_exp_forward, _so3_Jl, _so3_Jl_inv)) +def test_so3_small_angle_forward_and_gradients_are_finite(function): + def weighted_output(input): + output = function(input) + weights = torch.arange( + 1, output.numel() + 1, device=output.device, dtype=output.dtype + ).reshape(output.shape) + return output, (output * weights).sum() + + eager_input = _zero_and_small_angles(torch.float64).requires_grad_() + eager_output, eager_loss = weighted_output(eager_input) + (eager_gradient,) = torch.autograd.grad(eager_loss, eager_input) + + compiled = torch.compile(weighted_output, backend="eager", fullgraph=True) + compiled_input = eager_input.detach().clone().requires_grad_() + compiled_output, compiled_loss = compiled(compiled_input) + (compiled_gradient,) = torch.autograd.grad(compiled_loss, compiled_input) + + assert torch.isfinite(eager_output).all() + assert torch.isfinite(eager_gradient).all() + torch.testing.assert_close(compiled_output, eager_output) + torch.testing.assert_close(compiled_gradient, eager_gradient) + + def test_lietensor_tensor_alias_is_fullgraph_traceable_and_differentiable(): install_pypose_torch_compile_monkeypatch() data = pp.randn_SE3(4, sigma=0.2, dtype=torch.float64).tensor().detach().requires_grad_() From 1f1772c52ec81193f163fe156a6e8e72be0eb361 Mon Sep 17 00:00:00 2001 From: zhfang Date: Mon, 20 Jul 2026 21:28:09 +0000 Subject: [PATCH 10/10] fix comment: `__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. --- bae/utils/pypose_compile.py | 8 ++++++-- tests/autograd/test_pypose_compile.py | 23 +++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/bae/utils/pypose_compile.py b/bae/utils/pypose_compile.py index 9347cc7..0702980 100644 --- a/bae/utils/pypose_compile.py +++ b/bae/utils/pypose_compile.py @@ -80,8 +80,12 @@ def _dynamo_compatible_torch_function(cls, func, types, args=(), kwargs=None): # 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: - flat_args, _ = tree_flatten(args) - ltype = next(arg.ltype for arg in flat_args if isinstance(arg, lt.LieTensor)) + flat_inputs, _ = tree_flatten((args, kwargs)) + ltype = next( + arg.ltype + for arg in flat_inputs + if isinstance(arg, lt.LieTensor) + ) def wrap(tensor): if isinstance(tensor, torch.Tensor) and not isinstance(tensor, cls): diff --git a/tests/autograd/test_pypose_compile.py b/tests/autograd/test_pypose_compile.py index e5ee3e2..a2933c4 100644 --- a/tests/autograd/test_pypose_compile.py +++ b/tests/autograd/test_pypose_compile.py @@ -66,6 +66,29 @@ def test_lietensor_tensor_alias_is_fullgraph_traceable_and_differentiable(): assert inverse.ltype is pp.SE3_type +@pytest.mark.parametrize( + "operation", + ( + lambda value: torch.clone(input=value), + lambda value: torch.cat(tensors=(value, value), dim=0), + lambda value: torch.reshape(input=value, shape=(1, 2, 4)), + ), + ids=("clone", "cat", "reshape"), +) +def test_lietensor_keyword_inputs_preserve_ltype_in_fullgraph(operation): + value = pp.randn_SO3(2, dtype=torch.float64) + + expected = operation(value) + compiled = torch.compile(operation, backend="eager", fullgraph=True) + actual = compiled(value) + + assert isinstance(expected, pp.LieTensor) + assert isinstance(actual, pp.LieTensor) + assert expected.ltype is pp.SO3_type + assert actual.ltype is pp.SO3_type + torch.testing.assert_close(actual.tensor(), expected.tensor()) + + def test_ambient_se3_pipeline_is_fullgraph_traceable_with_local_gradients(): install_pypose_ambient_grad_monkeypatch() dtype = torch.float64