diff --git a/README.md b/README.md index 2a54b6f..381d40a 100644 --- a/README.md +++ b/README.md @@ -159,66 +159,6 @@ 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): diff --git a/doc/compile_lietensor.md b/doc/compile_lietensor.md new file mode 100644 index 0000000..fac2862 --- /dev/null +++ b/doc/compile_lietensor.md @@ -0,0 +1,59 @@ +### `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.