Jacrev vs autograd

Wesley_Neill: <p>I have an image that I want to take the gradient of at each pixel. I am able to get the gradient at every point on the image using a wrapper around <code>functional.grid_sample</code> <code>torch.autograd.grad</code>, and now I want to compare that output to the output of <code>functorch.jacrev</code>.</p>
<p>My wrapper works fine (assuming you are okay with garbage at the borders of each dimension), but using jacrev throws an error (entire stack trace at bottom). Any help in sorting this would be appreciated!</p>
<h3><a name="to-reproduce-1" class="anchor" href="#to-reproduce-1"></a>To reproduce:</h3>
<p>I’m using torch 1.11.0+cu113 on Ubuntu.</p>
<p>Run this code in a notebook if you want to reproduce it:</p>
<pre><code class="lang-auto">import matplotlib.pyplot as plt
import numpy as np
import torch
import functorch as ft
from torch.nn import functional as F


def image_func_gen(image): 
    """(I, C, H, W)"""

    def _normalize_coordinates(row_col_coords):
        """zero based, (P, 2 [row, col])"""
        
        image_dims = torch.tensor(image.shape)[-2:]
        ratio_spanned = (row_col_coords) / (image_dims - 1)
        scaled = ratio_spanned * 2 
        normalized_coords = -1 + scaled
        return normalized_coords
        
    def image_func(coords):
        """(P, 2 [row, col])"""
        
        normalized_coords = _normalize_coordinates(coords)
        normalized_coords = normalized_coords.flip(dims=(-1,))
        grid = normalized_coords[None, :, None, :]
        return F.grid_sample(input=image, grid=grid, align_corners=True)

    return image_func

def image_gradient(coords, image):
   
    # gradient at edges (last row/col) will be set to zero
    padded_image = torch.nn.functional.pad(image, (0, 1, 0, 1), mode='replicate')
    
    image_func = image_func_gen(padded_image)
    pixel_values = image_func(coords).squeeze()

    gradients = torch.autograd.grad(pixel_values, coords, torch.ones(len(pixel_values)))[0]
    return gradients

# Set up a simple image with identical gradient at every point:
height = 5.
width = 5.

h, w = int(height), int(width)
image = torch.arange(0, width).expand(h, w)

# create function to model image
image_function = image_func_gen(image[None, None])

# sample the function. output should be identical to original image. 
x_range = torch.linspace(0., width - 1, w, requires_grad=True)
y_range = torch.linspace(0., height - 1, h, requires_grad=True)
coords = torch.cartesian_prod(y_range, x_range) 
output = image_function(coords)

# now take the gradient:
gradient = image_gradient(coords, image[None, None])

# and compare to the Jacobian: (ERROR HERE)
ft.jacrev(image_function, argnums=0)(coords)
</code></pre>
<p>Results in:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c162a93062157ffbf4a3dd0b60c5634f8fd4046.png" data-download-href="https://discuss.pytorch.org/uploads/default/0c162a93062157ffbf4a3dd0b60c5634f8fd4046" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c162a93062157ffbf4a3dd0b60c5634f8fd4046.png" alt="image" data-base62-sha1="1IVdGvKEYPNuyuoUDRZuzmpEZVQ" width="690" height="312" data-dominant-color="F6F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1629×738 31.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="Wesley_Neill" data-post="1" data-topic="152019">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wesley_neill/48/25709_2.png" class="avatar"> Wesley_Neill:</div>
<blockquote>
<pre><code class="lang-auto">    gradients = torch.autograd.grad(pixel_values, coords, torch.ones(len(pixel_values)))[0]
</code></pre>
</blockquote>
</aside>
<p>If you’re trying to take the gradient of <code>gradients</code>, shouldn’t you specify <code>retain_graph=True</code> and <code>create_graph=True</code> when calling <code>torch.autograd.grad</code>?</p>
<p>If <code>retain_graph = False</code> (which is the default behavior) it won’t have a gradient and might result in an undefined Tensor?</p>
<p>Also, I’ve just noticed  you’re using <code>jacrev</code> which is from <code>functorch</code> and not pytorch, that’s actually a functional transformation and I’m not sure how that interaction with standard pytorch gradient ops like <code>torch.autograd.grad</code>, it might be best opening an issue on <a href="https://github.com/pytorch/functorch/issues" class="inline-onebox" rel="noopener nofollow ugc">Issues · pytorch/functorch · GitHub</a> as this seems more of functorch issue than pytorch!</p>