Using backward on matrix custom function

Vendrick17: <p>Hello everyone, I am implementing QR decomposition using modified Gram-Schmidt in a function and I am trying to get the gradient, but I am struggling, my function is:</p>
<pre><code class="lang-auto">def modified_gram_schmidt(A):

    m, n = A.shape
    Q = torch.zeros((m, n))
    R = torch.zeros((n, n))
    
    for j in range(n):
        v = A[:, j]
        for i in range(j):                 
            R[i, j] = torch.dot(Q[:, i], v)
            v = v - R[i, j] * Q[:, i]
        R[j, j] = torch.norm(v)
        Q[:, j] = v / R[j, j]

    return Q, R
</code></pre>
<p>I am trying to compute its gradient using <code>backward()</code> as follows:</p>
<pre><code class="lang-auto">Q, R = modified_gram_schmidt(A)

Q.sum().backward()
dQ1_dA = A.grad
print(dQ1_dA)
</code></pre>
<p>And it outputs <code>None</code> as result, I think that this because the operations are not being tracked by <code>Autograd</code>, I created <code>A</code> with <code>requireds_grad=True</code>, any suggestion or idea to compute it correctly will be appreciated.</p>

soulitzer: <p>Does the following work for you?</p>
<pre><code class="lang-python">import torch

def modified_gram_schmidt(A):
    m, n = A.shape
    # Set requires grad so that gradients propagate through, clone to avoid
    # in-placing over a leaf tensor
    Q = torch.zeros((m, n), requires_grad=True).clone()
    R = torch.zeros((n, n), requires_grad=True).clone()

    for j in range(n):
        v = A[:, j]
        for i in range(j):
            R[i, j] = torch.dot(Q[:, i], v)
            v = v - R[i, j] * Q[:, i]
        R[j, j] = torch.norm(v)
        Q[:, j] = v / R[j, j]

    return Q, R

A = torch.rand(3, 5, requires_grad=True)

# This context manager is available in versions &gt;=2.0.0
with torch.autograd.graph.allow_mutation_on_saved_tensors():
  Q, R = modified_gram_schmidt(A)
  Q.sum().backward()

dQ1_dA = A.grad
print(dQ1_dA)
</code></pre>

Vendrick17: <p>Thank you so much! It worked but the result doesn’t coincide with the one I got from taking the gradient on the <code>torch.linalg.qr(A)</code>, I did it as follows:</p>
<pre><code class="lang-auto">Q1, R1 = torch.linalg.qr(A)# torch version

loss1 = Q1.sum()
loss1.backward()
dQ1_dA = A.grad
print(dQ1_dA)
A.grad.zero_()
</code></pre>
<p>I also set the random seed to the same number and ensure the matrix A is the same for functions, probably my way to get this last grad has an issue.</p>

soulitzer: <p>Does forward produce the same result? Is it possible that the function is poorly conditioned on certain inputs? Have you tried a simple square matrix example?</p>
<p>Checking the gradients with gradcheck, it seems fine:</p>
<pre><code class="lang-python">A = torch.tensor([
  [1., 2., 3.],
  [5., 3., 2.],
  [4., 1., 7.]
], requires_grad=True, dtype=torch.float64)

with torch.autograd.graph.allow_mutation_on_saved_tensors():
  torch.autograd.gradcheck(modified_gram_schmidt, A)
</code></pre>

Vendrick17: <p>You are right, the forward pass of my function doesn’t provide the same exact result for <code>Q</code> and <code>R</code> (Varies in some signs), but it passes the test when I compute <code>Q@R</code> and compares it with the product obtained from <code>torch</code> library, but it is difficult to get the same exact result provided by torch since their method is probably different.</p>
<p>And what do you mean by poorly condtioned?</p>
<p>I also ran the last code you sent and there was no errors.</p>

soulitzer: <p>Oh I mean ill-conditioned - the derivative at certain inputs may be steep which can cause the gradient computed analytically to differ significantly from the one computed numerically via finite differencing.</p>