How does torch.where work in autograd

wasabi_linguist: <p>Hi, sorry if the answer is trivial but I can’t seem to get how <code>autograd</code> differentiates <code>torch.where</code> operations where the mask used by <code>torch.where</code> is also derived from an input.</p>

Arunprakash-A: <p>Hi</p>
<p>That’s a good question.</p>
<p>Let’s walk through three cases to see how gradients behave with <code>torch.where()</code>.<br>
<strong>Reminder:</strong> gradients in PyTorch flow only to <em>differentiable inputs</em> (like model weights or tensors with <code>requires_grad=True</code>). Here, we want to know how <code>torch.where</code> routes the gradient</p>
<p>Suppose we have</p>
<pre data-code-wrap="python"><code class="lang-python">x = [1,2,3,4]
y = [-1,-2,3,4]
</code></pre>
<h1><a name="p-474949-case-1-constant-1" class="anchor" href="#p-474949-case-1-constant-1"></a>Case-1: Constant</h1>
<p><code>z = torch.where(x&gt;2,1,0)</code></p>
<p><strong>output:</strong>[0,0,1,1].<br>
Here, the result is constant with respect to both <code>x</code> and <code>y</code>. There’s <strong>no functional relationship</strong> between inputs and output, so <strong>no gradient can flow</strong></p>
<h1><a name="p-474949-case-2-function-of-x-2" class="anchor" href="#p-474949-case-2-function-of-x-2"></a>Case-2: Function of x</h1>
<p><code>z = torch.where(x&gt;2,x,0)</code><br>
<strong>output:</strong>[0,0,3,4].</p>
<p>Now, if the condition is <code>True</code>, the corresponding element of <code>x</code> is selected.</p>
<ul>
<li>For those positions, gradients flow back to <code>x</code></li>
<li>For the others (masked out), the gradient is zero</li>
</ul>
<p>So backpropagation works, but only through the “active” elements of <code>x</code></p>
<h1><a name="p-474949-case-3-function-of-both-x-and-y-3" class="anchor" href="#p-474949-case-3-function-of-both-x-and-y-3"></a>Case-3: function of both x and y</h1>
<p><code>z = torch.where(x&gt;2,x,y)</code><br>
<strong>output:</strong> will be [-1,-2,3,4].</p>
<p>Here:</p>
<ul>
<li>Where the condition is <code>True</code>, the gradient flows into <code>x</code></li>
<li>Where the condition is <code>False</code>, the gradient flows into <code>y</code></li>
</ul>
<p>Thus, both <code>x</code> and <code>y</code> can receive gradients, depending on the mask.</p>
<p>Mathematically, the operator is <strong>deterministic</strong>, therefore, we can compute the gradient.To my knowledge, only <strong>stochastic</strong> operators (like sampling) break down the flow of gradients!</p>
<p>Hope it helps!</p>

valerian.rey: <p>Sorry if I misunderstand your question, but are you asking how torch can differentiate the result of <code>torch.where</code> with respect to its <code>condition</code> parameter? It doesn’t: <code>condition</code> is a <code>BoolTensor</code>, that will never require grad. So it’s considered as a constant.</p>
<p>However, as explained by <a class="mention" href="/u/arunprakash-a">@Arunprakash-A</a>, torch can differentiate the result of <code>torch.where</code> with respect to either <code>input</code> and / or <code>other</code> (depending on their <code>requires_grad</code> field). And this differentiation depends on the value of <code>condition</code>.</p>