Suggestions for making this equation autograd-able?

LOGUNIVPM: <p>Hi, I am dealing with a spintronic problem where the loss is a Hamiltonian of the type:</p>
<pre><code class="lang-auto">        energy = 0
        for i in range(self.Npbits):
            tmp = 0.5 * torch.dot(self.J[i], pbits) + self.h[i]
            energy -= tmp * pbits[i]
        return energy
</code></pre>
<p>J is a square matrix, and J[i] one of its lines. pbits is a vector of the same length.<br>
I doubt this solution will not allow the gradient to be automatically computed because of the “for” loop and because I’m taking the rows of J by indexing. Am I wrong? Is there a definitive guide to writing custom code that can be autograd-able?</p>
<p>Best</p>

ptrblck: <p>The loop and the indexing operation won’t detach the tensor from the computation graph. Did you see any issues using your approach?</p>

KFrank: <p>Hi Leonardo!</p>
<aside class="quote no-group" data-username="LOGUNIVPM" data-post="1" data-topic="207261" data-full="true">
<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/logunivpm/48/65933_2.png" class="avatar"> LOGUNIVPM:</div>
<blockquote>
<pre><code class="lang-auto">        energy = 0
        for i in range(self.Npbits):
            tmp = 0.5 * torch.dot(self.J[i], pbits) + self.h[i]
            energy -= tmp * pbits[i]
        return energy
</code></pre>
</blockquote>
</aside>
<p>As an aside, you can rewrite this without the loop:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.4.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt;
&gt;&gt;&gt; Npbits = 5
&gt;&gt;&gt; J = torch.randn (Npbits, Npbits)
&gt;&gt;&gt; pbits = torch.randn (Npbits)
&gt;&gt;&gt; h = torch.randn (Npbits)
&gt;&gt;&gt;
&gt;&gt;&gt; energy = 0   # compute with loop
&gt;&gt;&gt; for  i in range (Npbits):
...     tmp = 0.5 * torch.dot (J[i], pbits) + h[i]
...     # print (tmp.shape)
...     energy -= tmp * pbits[i]
...
&gt;&gt;&gt; # loop-free version
&gt;&gt;&gt; energyB = -(0.5 * torch.einsum ('ij, i, j -&gt;', J, pbits, pbits) + torch.einsum ('i, i -&gt;', h, pbits))
&gt;&gt;&gt;
&gt;&gt;&gt; energy
tensor(4.5874)
&gt;&gt;&gt; energyB
tensor(4.5874)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

LOGUNIVPM: <p>Thank you! Indeed I was approaching that kind of solution with matmul e dot, I didn’t know einsum, very useful…<br>
The solution works and the network backpropagates effectively.</p>
<p>Since you’ve been so kind, I’m going to bother you with another similar problem <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"><br>
This time the gradient is harder to compute, since it uses a sign function (which I replaced with a tanh and a strong multiplying factor) and conditionals inside a for loop.</p>
<p>The original code is the following:</p>
<pre><code>    loss = 0
    pbits_variables = self.pbits[self.index_variables]
    for i, clause in enumerate(self.clauses):
        truth_value = 0
        for variable in clause:
            if pbits_variables[np.abs(variable) - 1] * np.sign(variable) &gt; 0:
                truth_value = 1
                break
        self.loss += (1 - truth_value)
</code></pre>
<p>clauses are a list of rows of numbers (inside a list), therefore for each number in a clause the if checks whether the &gt; 0 condition is verified, and in that case, it breaks and says that truth_value = 1. I can do the same checking the condition on the array and using clamp, so that the truth_value is 1 if at least one element in the array is &gt; 0.</p>
<p>I have devised this code:</p>
<pre><code>    pbits_variables = pbits[:,self.index_variables]
    cost = torch.zeros(1, requires_grad=True)
    for i, c in enumerate(self.clauses):
        truth_value = torch.zeros(1, requires_grad=True)
        clause = torch.from_numpy(np.array(c))
        v = pbits_variables[:,torch.abs(clause) - 1] * nn.functional.tanh(DENORM_BETA*clause)
        truth_value = torch.clamp(torch.sum(v &gt; 0), max=1)
        cost += (1 - truth_value)
</code></pre>
<p>I’ve probably missed something since the last line bumps this error message.</p>
<pre><code>  RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.
</code></pre>

AlphaBetaGamma96: <aside class="quote no-group" data-username="LOGUNIVPM" data-post="4" data-topic="207261">
<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/logunivpm/48/65933_2.png" class="avatar"> LOGUNIVPM:</div>
<blockquote>
<pre><code class="lang-auto">        cost += (1 - truth_value)
</code></pre>
</blockquote>
</aside>
<p>You shouldn’t have any in-place operations within your function, so remove any <code>+=</code> or <code>*=</code> etc. and see if that resolves the issue.</p>

LOGUNIVPM: <p>You are right, just using cost = cost + … instead of += is the way to go.<br>
I am facing convergence issues but that is probably related to the nature of the problem. I can see a nonzero gradient flow through the network, therefore the loss as I coded it should be able to compute the gradient automatically. This seems to me quite sound… (PS: the plot is done according to this old discussion <a href="https://discuss.pytorch.org/t/check-gradient-flow-in-network/15063/7" class="inline-onebox">Check gradient flow in network - #7 by RoshanRane</a>)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc535eec96eaa18d74dba88bf11f5d38177cd2d4.png" data-download-href="https://discuss.pytorch.org/uploads/default/dc535eec96eaa18d74dba88bf11f5d38177cd2d4" title="grad_flow"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc535eec96eaa18d74dba88bf11f5d38177cd2d4_2_628x500.png" alt="grad_flow" data-base62-sha1="vr5FcKSPoiMZ8BNANzjkfdVSsmg" width="628" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc535eec96eaa18d74dba88bf11f5d38177cd2d4_2_628x500.png, https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc535eec96eaa18d74dba88bf11f5d38177cd2d4.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc535eec96eaa18d74dba88bf11f5d38177cd2d4.png 2x" data-dominant-color="F3F3F9"><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">grad_flow</span><span class="informations">640×509 43.4 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>