Repository: pytorch
Issue: Gradient of image rotation
Labels: triaged,enhancement,module: forward ad

### 🚀 The feature, motivation and pitch

My goal is to have a function that depends on an angle and to compute the gradient with respect to the rotation angle. Here is a MWE of the code I'm trying to write:

```
        def get_img_pixel(angle):
            return kornia.geometry.transform.rotate(torch.randn(1,1,10,10), angle)[0,0,6,6]
        
        torch.func.jvp(get_img_pixel, (torch.zeros(1),), (torch.ones(1),))
```
I get the following stack trace, which mentions to submit an issue here.

```----------------------------------------------------------
NotImplementedError      Traceback (most recent call last)
Cell In[60], line 4
      1 def get_img_pixel(angle):
      2     return kornia.geometry.transform.rotate(torch.randn(1,1,10,10), angle)[0,0,6,6]
----> 4 torch.func.jvp(get_img_pixel, (torch.zeros(1),), (torch.ones(1),))

File /state/partition1/llgrid/pkg/anaconda/python-LLM-2023b/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py:916, in jvp(func, primals, tangents, strict, has_aux)
    864 @exposed_in("torch.func")
    865 def jvp(func: Callable, primals: Any, tangents: Any, *, strict: bool = False, has_aux: bool = False):
    866     """
    867     Standing for the Jacobian-vector product, returns a tuple containing
    868     the output of `func(*primals)` and the "Jacobian of ``func`` evaluated at
   (...)
    913 
    914     """
--> 916     return _jvp_with_argnums(func, primals, tangents, argnums=None, strict=strict, has_aux=has_aux)

File /state/partition1/llgrid/pkg/anaconda/python-LLM-2023b/lib/python3.10/site-packages/torch/_functorch/vmap.py:39, in doesnt_support_saved_tensors_hooks.<locals>.fn(*args, **kwargs)
     36 @functools.wraps(f)
     37 def fn(*args, **kwargs):
     38     with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39         return f(*args, **kwargs)

File /state/partition1/llgrid/pkg/anaconda/python-LLM-2023b/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py:965, in _jvp_with_argnums(func, primals, tangents, argnums, strict, has_aux)
    963     primals = _wrap_all_tensors(primals, level)
    964     duals = _replace_args(primals, duals, argnums)
--> 965 result_duals = func(*duals)
    966 if has_aux:
    967     if not (isinstance(result_duals, tuple) and len(result_duals) == 2):

Cell In[60], line 2, in get_img_pixel(angle)
      1 def get_img_pixel(angle):
----> 2     return kornia.geometry.transform.rotate(torch.randn(1,1,10,10), angle)[0,0,6,6]

File ~/.local/lib/python3.10/site-packages/kornia/geometry/transform/affwarp.py:287, in rotate(tensor, angle, center, mode, padding_mode, align_corners)
    284 rotation_matrix: torch.Tensor = _compute_rotation_matrix(angle, center)
    286 # warp using the affine transform
--> 287 return affine(tensor, rotation_matrix[..., :2, :3], mode, padding_mode, align_corners)

File ~/.local/lib/python3.10/site-packages/kornia/geometry/transform/affwarp.py:161, in affine(tensor, matrix, mode, padding_mode, align_corners)
    159 height: int = tensor.shape[-2]
    160 width: int = tensor.shape[-1]
--> 161 warped: torch.Tensor = warp_affine(tensor, matrix, (height, width), mode, padding_mode, align_corners)
    163 # return in the original shape
    164 if is_unbatched:

File ~/.local/lib/python3.10/site-packages/kornia/geometry/transform/imgwarp.py:206, in warp_affine(src, M, dsize, mode, padding_mode, align_corners, fill_value)
    203 # src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
    204 src_norm_trans_dst_norm = _torch_inverse_cast(dst_norm_trans_src_norm)
--> 206 grid = F.affine_grid(src_norm_trans_dst_norm[:, :2, :], [B, C, dsize[0], dsize[1]], align_corners=align_corners)
    208 if padding_mode == "fill":
    209     return _fill_and_warp(src, grid, align_corners=align_corners, mode=mode, fill_value=fill_value)

File /state/partition1/llgrid/pkg/anaconda/python-LLM-2023b/lib/python3.10/site-packages/torch/nn/functional.py:4341, in affine_grid(theta, size, align_corners)
   4338 elif min(size) <= 0:
   4339     raise ValueError("Expected non-zero, positive output size. Got {}".format(size))
-> 4341 return torch.affine_grid_generator(theta, size, align_corners)

NotImplementedError: Trying to use forward AD with affine_grid_generator that does not support it because it has not been implemented yet.
Please file an issue to PyTorch at https://github.com/pytorch/pytorch/issues/new?template=feature-request.yml so that we can prioritize its implementation.
Note that forward AD support for some operators require PyTorch to be built with TorchScript and for JIT to be enabled. If the environment var PYTORCH_JIT=0 is set or if the library is not built with TorchScript, some operators may no longer be used with forward AD.```

Thanks for your help in advance.

### Alternatives

If there's a simpler way to do this, please let me know! Performance is not big issue yet, and I'm open to temporary approaches.

### Additional context

_No response_