pypose.autograd.function.parallel_for_sparse_jacobian

pypose.autograd.function.parallel_for_sparse_jacobian(function)

The parallel_for_sparse_jacobian() (alias psjac()) decorator allows PyPose’s optional sparse backend to trace and assemble sparse Jacobians more efficiently. It wraps a batched function whose batch samples are independent, so the function can be treated as a parallel computation along the batch dimension during sparse Jacobian construction. Here, the leading dimension of each input and output tensor is the batch dimension.

Warning

This decorator is required by sparse pypose.optim.LevenbergMarquardt with sparse=True.

Note

See runable example for bundle adjustment.

Example

@psjac
def edge_error(node1, node2, relative):
    # node1: pp.SE3 (N, 7), node2: pp.SE3 (N, 7), relative: pp.SE3 (N, 7)
    # returns: (N, 6)
    return (relative.Inv() @ node1.Inv() @ node2).Log().tensor()

Here, each output is the error for one pose-graph edge, computed only from the matching input of node1, node2, and relative.

Warning

Similar to torch.vmap, it should not be used for functions that mix information across rows, such as batch reductions, global statistics, or any computation where one output in the batch depends on multiple input batch samples.

Note

This decorator doesn’t change the function behavior. It only adds tracing information for sparse Jacobian construction.