How do I pass grad through torch.where

username: themoonboy
cooked: <p>Hi guys, when using <em><strong>torch.where</strong></em>, it works as a threshold thus unable to pass the grad by torch.autograd.grad(…). Instead, is there a way to keep to same graph and pass through to next differentiable operation?<br>
For example, in a very simple code below</p>
<pre><code class="lang-auto">import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
b = torch.tensor([[0.1],[0.9]])
c = torch.mm(a, b)

# ** position2 **: grad couldn't pass back here
c = torch.where(c&gt;2, torch.tensor(1.0, requires_grad=True), torch.tensor(0.0, requires_grad=True))
# ** position1 **: the most place grad could arrive
loss = c.sum()

print(torch.autograd.grad(loss, a, allow_unused=True))
</code></pre>
<p>The <strong>c</strong> uses torch.where as a threshold to get binary values, of course the grad cannot arrive at <strong>position2</strong>. I just wonder is there some methods that we can keep c’s grad in <strong>position1</strong> and directly paste to <strong>position2</strong>, then keep calculating grads until arrive to <strong>a</strong>.</p>
<p>Above is only a hand-written example, when using build-in optimizer, after calculating loss and doing loss.backward(), is there any methods to let the backward <strong>overpass</strong> torch.where() like example above?</p>

username: albanD
cooked: <p>Hi,</p>
<p>You can. write a custom Function whose backward will be the identity while the forward is your torch.where.<br>
You can find <a href="https://pytorch.org/docs/stable/notes/extending.html">here</a> how to write such Function.</p>

username: themoonboy
cooked: <p>Problem solved! Thanks so much</p>

username: stijnvdl
cooked: <p>I’m struggling with the implementation.<br>
This is my current code (I am using PyTorch 1.12.1):</p>
<pre><code class="lang-auto">class CustomWhere(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        result = torch.where(input &gt; 0.1, 1.0, 0.0)
        return result
        
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return torch.ones(input.shape[0])

def custom_where(x):
    result = CustomWhere.apply(x)
    return result
</code></pre>
<p>which I think implements what you mentioned. I’ve also tried</p>
<pre><code class="lang-auto">return grad_output * torch.ones(input.shape[0])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">return torch.eye(input.shape[0])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">return grad_output * torch.eye(input.shape[0])
</code></pre>
<p>I test all the approaches using gradcheck</p>
<pre><code class="lang-auto">from torch.autograd import gradcheck
input = torch.randn(20, dtype=torch.double,requires_grad=True)
test = gradcheck(custom_where, input, eps=1e-6, atol=1e-4)
</code></pre>
<p>but they all unfortunately fail with</p>
<blockquote>
<p>GradcheckError: Jacobian mismatch for output 0 with respect to input 0</p>
</blockquote>
<p>What am I doing wrong?</p>

username: albanD
cooked: <p>Hey!</p>
<p>Your backward here is not an identity, it returns a constant 1 <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
simply do <code>return grad_output</code> in there.<br>
This is because the custom Function compute a vector jacobian product step. Not the jacobian itself.</p>