pypose.func.jacrev

pypose.func.jacrev(func, argnums=0, *, has_aux=False, chunk_size=None, _preallocate_and_copy=False)

This function provides the exact same functionality as torch.func.jacrev(), except that it allows LieTensor to be used as input when calculating the jacobian.

Parameters:

• func (function) – A Python function that takes one or more arguments, one of which must be a Tensor, and returns one or more Tensors

• argnums (int or Tuple[int]) – Optional, integer or tuple of integers, saying which arguments to get the Jacobian with respect to. Default: 0.

• has_aux (bool) – Flag indicating that func returns a (output, aux) tuple where the first element is the output of the function to be differentiated and the second element is auxiliary objects that will not be differentiated. Default: False.

• chunk_size (None or int) – If None (default), use the maximum chunk size (equivalent to doing a single vmap over vjp to compute the jacobian). If 1, then compute the jacobian row-by-row with a for-loop. If not None, then compute the jacobian chunk_size rows at a time (equivalent to doing multiple vmap over vjp). If you run into memory issues computing the jacobian, please try to specify a non-None chunk_size.

Returns:

Returns a function that takes in the same inputs as func and returns the Jacobian of func with respect to the arg(s) at argnums. If has_aux is True, then the returned function instead returns a (jacobian, aux) tuple where jacobian is the Jacobian and aux is auxiliary objects returned by func.

A basic usage with our LieTensor type would be the transformation function.

>>> import pypose as pp
>>> import torch
>>> def func(pose, points):
...     return pose @ points
>>> pose = pp.randn_SE3(1)
>>> points = torch.randn(1, 3)
>>> jacobian = pp.func.jacrev(func)(pose, points)
>>> jacobian
tensor([[[[ 1.0000,  0.0000,  0.0000,  0.0000,  1.5874, -0.2061,  0.0000]],
        [[ 0.0000,  1.0000,  0.0000, -1.5874,  0.0000, -1.4273,  0.0000]],
        [[ 0.0000,  0.0000,  1.0000,  0.2061,  1.4273,  0.0000,  0.0000]]]])