How to implement a piecewise function that supports autograd?

username: zyzhu2000
cooked: <p>I want to implement a piecewise function on a tensor that has the following logic for each element in the tensor.</p>
<pre><code class="lang-auto">def f(x):
    if x==0.5:
        return 1
    else:
        return torch.atanh(x)/(x-0.5)
</code></pre>
<p>Clearly, torch.atanh(x)/(x-0.5) is not defined when x==0.5.</p>
<p>The above function is for one single element, and so I can easily make two branches. But I want to know how to implement this logic on a multi-element tensor X, such that a tensor of the same shape as the input is returned and each element has the value f(x), where x is the corresponding element in X.</p>
<p>I could use the torch.where() function, but then since both branches of torch.where() are computed, I may not avoid divide by zero errors when the elelemtn is 0.5.</p>

username: soulitzer
cooked: <p><code>torch.where</code> should be okay because if you divide by 0, shouldn’t you be getting inf/nan?</p>

username: AlphaBetaGamma96
cooked: <p>This might be a potential solution to your question,</p>
<pre><code class="lang-auto">def fnew(x):
  return (~(x==0.5))*torch.nan_to_num(torch.atanh(x)/(x-0.5)) + (x==0.5)*1
</code></pre>
<p>This would be a branchless example of your if-else statement above. If <code>x != 0.5</code> (as defined by <code>(~(x==0.5))</code> ) the left hand of the is multipled by 0 and it returns 1. If <code>x == 0.5</code>, the opposite behaviour happens and it returns <code>atanh(x)/(x-0.5)</code>.</p>
<p>I’ve only placed a <code>torch.nan_to_num</code> in there so to convert any <code>nans</code> to 0. If your inputs are only within the range <code>(-1 ,1)</code> you can remove the <code>torch.nan_to_num</code> as it’ll serve no purpose.</p>

username: Gordon_research
cooked: <p>is autograd able to handle this piecewise function? Is it differentiating on each sub-interval?</p>

username: Gordon_research
cooked: <p>Looks like it’s able to handle:</p>
<p>import torch<br>
def compute_krw(Sw):<br>
return Sw**3</p>
<p>Sw = torch.rand(4096, 50).requires_grad_(True)<br>
krw = compute_krw(Sw)<br>
krw.backward(torch.ones(krw.shape))</p>
<p>print(Sw)<br>
tensor([[0.7071, 0.3658, 0.8739,  …, 0.1909, 0.4205, 0.8919],<br>
[0.8442, 0.5447, 0.8630,  …, 0.7269, 0.8380, 0.0856],<br>
[0.1522, 0.3359, 0.1544,  …, 0.2483, 0.9337, 0.4175],<br>
…,<br>
[0.6864, 0.4651, 0.4591,  …, 0.9432, 0.3384, 0.3128],<br>
[0.8458, 0.0655, 0.3948,  …, 0.5267, 0.7302, 0.8763],<br>
[0.8550, 0.8541, 0.4046,  …, 0.9442, 0.1375, 0.4865]],<br>
requires_grad=True)</p>
<p>print(Sw.grad)<br>
tensor([[ 4.1874e+00,  6.9268e-01,  0.0000e+00,  …,  2.8091e-02,<br>
1.0571e+00,  0.0000e+00],<br>
[ 0.0000e+00,  2.1619e+00,  0.0000e+00,  …,  4.4784e+00,<br>
0.0000e+00, -3.8395e-03],<br>
[-1.5801e-02,  5.2467e-01, -1.4277e-02,  …,  1.6212e-01,<br>
0.0000e+00,  1.0353e+00],<br>
…,<br>
[ 3.8929e+00,  1.4097e+00,  1.3591e+00,  …,  0.0000e+00,<br>
5.3783e-01,  4.1034e-01],<br>
[ 0.0000e+00,  2.1532e-02,  8.7634e-01,  …,  1.9774e+00,<br>
4.5290e+00,  0.0000e+00],<br>
[ 0.0000e+00,  0.0000e+00,  9.4362e-01,  …,  0.0000e+00,<br>
-2.2675e-02,  1.5962e+00]])</p>

username: Jamie_Donnelly
cooked: <p>Autograd can handle control flow within models since the autograd graph is recreated after each <code>.backward()</code> call on the loss. The autograd graph can be different during each iteration as a result of control flow in model’s forward methods, but the graph is created as normal containing the backward methods for whichever path of the control flow get executed.</p>
<p>For the sake of performance though, <code>torch.where</code> operations will probably be faster than writing control flow in python.</p>

username: ulveea
cooked: <p>If the gradient graph is created on every iteration (i.e. on every backward pass of a batch), what happens if individual sample outputs end up in the domain of different sub-functions of the piecewise function? Surely, these samples should have different gradient graphs.</p>