Is ``torch.autograd.grad" differentiable for its internal variables?

username: 111414
cooked: <p>I have a network that incorporates a learnable gradient descent step and is implemented as follows:</p>
<hr>
<p><strong>Implementation <span class="hashtag">#1:</span></strong></p>
<pre><code class="lang-python">import torch
from torch import nn

def cal_grad(x, loss):
    v = torch.autograd.Variable(x, requires_grad=True)
    loss_v = loss(v)
    return torch.autograd.grad(loss_v, v, torch.ones_like(loss_v))[0]

class Net_autograd(nn.Module):
    def __init__(self):
        super(Net_autograd, self).__init__()
        self.A = nn.Parameter(torch.rand(100, 50))
        self.rho = nn.Parameter(torch.rand(1,))

    def forward(self, x_init, y):
        grad = cal_grad(x_init, lambda z: 0.5 * (z.mm(self.A) - y).pow(2).sum(dim=[1,]))  # calculate gradient implicitly (without giving its formula)
        x_hat = x_init - self.rho * grad
        return x_hat

net = Net_autograd()
x_init = torch.rand(10, 100, requires_grad=True)
y = torch.rand(10, 50, requires_grad=False)
x_gt = torch.rand(10, 100, requires_grad=False)

x_hat = net(x_init, y)

loss = (x_hat - x_gt).pow(2).mean()  # L2 loss
loss.backward()

print('A', net.A.grad)
print('rho', net.rho.grad)
print('x_init', x_init.grad)
</code></pre>
<hr>
<p>It is not that common that I use a learnable matrix <code>self.A</code> to calculate gradient and obtain <code>x_hat</code>, in its forwarding process. I am confused by that why after <code>loss.backward()</code>, the gradient of parameter <code>self.A</code> is <code>None</code>? <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>
<p>And I re-implement the original network by explicitly giving the analytic gradient calculation formula, the parameter <code>self.A</code> can be successfully learned since <code>net.A.grad</code> will not be <code>None</code>:</p>
<hr>
<p><strong>Implementation <span class="hashtag">#2:</span></strong></p>
<pre><code class="lang-python">import torch
from torch import nn

class Net_explicit(nn.Module):
    def __init__(self):
        super(Net_explicit, self).__init__()
        self.A = nn.Parameter(torch.rand(100, 50))
        self.rho = nn.Parameter(torch.rand(1,))

    def forward(self, x_init, y):
        grad = (x_init.mm(self.A) - y).mm(self.A.t())  # calculate gradient explicitly
        x_hat = x_init - self.rho * grad
        return x_hat

net = Net_explicit()
x_init = torch.rand(10, 100, requires_grad=True)
y = torch.rand(10, 50, requires_grad=False)
x_gt = torch.rand(10, 100, requires_grad=False)

x_hat = net(x_init, y)

loss = (x_hat - x_gt).pow(2).mean()  # L2 loss
loss.backward()

print('A', net.A.grad)
print('rho', net.rho.grad)
print('x_init', x_init.grad)
</code></pre>
<hr>
<p>I want to ask that:</p>
<ol>
<li>
<p>The second implementation works well as my expectation, but why the first one does not train <code>self.A</code> (with <code>None</code> gradient)? I guess that the <code>torch.autograd.grad</code> itself may not be differentiable for its internal variables, since the external learnable gradient descent step size <code>self.rho</code> and input <code>x_init</code> obtain their own <code>grad</code>s.</p>
</li>
<li>
<p>To implement more complicated pipelines, in which <code>self.A</code> may be even a large network, what can I do to make the implementation <strong><span class="hashtag">#1</span></strong> works well as <strong><span class="hashtag">#2</span></strong> (i.e., to make <code>self.A</code> be trainable with <code>autograd</code> mechanisms)? <img src="https://discuss.pytorch.org/images/emoji/apple/cold_sweat.png?v=12" title=":cold_sweat:" class="emoji" alt=":cold_sweat:" loading="lazy" width="20" height="20"></p>
</li>
<li>
<p>By the way, I am not sure if the <code>grad</code> of <code>x_init</code> calculated in the forwardings of <strong><span class="hashtag">#1</span></strong> and <strong><span class="hashtag">#2</span></strong> are identical or totally equal.  Could you please check this for me?</p>
</li>
</ol>
<p>Note: The first dimensions of <code>x_init</code>, <code>y</code> and <code>x_gt</code>, with the size of 10, are the batch dimensions. The version of my <code>torch</code> is <code>1.7.1</code>, and the <code>torch.autograd.grad</code> does not have the parameter <code>is_grads_batched</code>.</p>

username: 111414
cooked: <p>It seems that I solved this problem by modifying the returning code of function <code>cal_grad</code> to be:</p>
<pre data-code-wrap="Python"><code class="lang-nohighlight">    return torch.autograd.grad(loss_v, v, torch.ones_like(loss_v), create_graph=True, retain_graph=True)[0]
</code></pre>