<p dir="auto"><strong>Bug</strong><br>
<code class="notranslate">jnp.flatnonzero</code> and <code class="notranslate">jnp.argwhere</code> mention <code class="notranslate">fill_value</code> option in documentation, but calling with <code class="notranslate">fill_value</code> leads to 'unexpected keyword argument' error.</p>
<p dir="auto">In the below code, the last two lines will raise errors due to unexpected argument <code class="notranslate">fill_value</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
d = jnp.array([True,False,True])
jnp.nonzero(d, size=20, fill_value=100)
jnp.flatnonzero(d, size=20, fill_value=100)
jnp.argwhere(d, size=20, fill_value=100)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-s1">d</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">True</span>,<span class="pl-c1">False</span>,<span class="pl-c1">True</span>])
<span class="pl-s1">jnp</span>.<span class="pl-en">nonzero</span>(<span class="pl-s1">d</span>, <span class="pl-s1">size</span><span class="pl-c1">=</span><span class="pl-c1">20</span>, <span class="pl-s1">fill_value</span><span class="pl-c1">=</span><span class="pl-c1">100</span>)
<span class="pl-s1">jnp</span>.<span class="pl-en">flatnonzero</span>(<span class="pl-s1">d</span>, <span class="pl-s1">size</span><span class="pl-c1">=</span><span class="pl-c1">20</span>, <span class="pl-s1">fill_value</span><span class="pl-c1">=</span><span class="pl-c1">100</span>)
<span class="pl-s1">jnp</span>.<span class="pl-en">argwhere</span>(<span class="pl-s1">d</span>, <span class="pl-s1">size</span><span class="pl-c1">=</span><span class="pl-c1">20</span>, <span class="pl-s1">fill_value</span><span class="pl-c1">=</span><span class="pl-c1">100</span>)</pre></div>
<p dir="auto"><strong>Why?</strong><br>
The implementation of <code class="notranslate">jnp.flatnonzero</code> and <code class="notranslate">jnp.argwhere</code> link back to <code class="notranslate">jnp.nonzero</code> but the <code class="notranslate">fill_value</code> argument wasn't threaded through (just the <code class="notranslate">size</code> argument).</p>
<p dir="auto">The threading is there in <code class="notranslate">jnp.where</code> and other cases, just those two have been missed as far as I've noticed.</p>
<p dir="auto"><strong>Fix</strong><br>
Pull request:<br>
<a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1078666664" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/8926" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/8926/hovercard" href="https://github.com/google/jax/pull/8926">#8926</a><br>
just threads the argument</p>
<p dir="auto">Images</p>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/44170519/145838724-958be84b-68b3-4d8d-980e-4dc8bc5bb0d3.png"><img src="https://user-images.githubusercontent.com/44170519/145838724-958be84b-68b3-4d8d-980e-4dc8bc5bb0d3.png" alt="image" style="max-width: 100%;"></a></p>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/44170519/145838907-673fb312-8b0b-4433-9be3-0d29cdf5cd82.png"><img src="https://user-images.githubusercontent.com/44170519/145838907-673fb312-8b0b-4433-9be3-0d29cdf5cd82.png" alt="image" style="max-width: 100%;"></a></p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
</ul>
<p dir="auto">I'm not sure if there's an elegant way to fix this. I noticed that when I differentiate <code class="notranslate">jnp.linalg.norm</code> when the input is <code class="notranslate">0</code>, the gradient is <code class="notranslate">nan</code>. I believe this is due to the vanishing denominator in the derivative. However, the numerator compensates for this in reality and the gradient should be <code class="notranslate">0</code>. For example,</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp

x = jnp.array([1., 2., 3.])
jax.grad(jnp.linalg.norm)(x - x) #[nan nan nan]

l2_dist_sq = lambda x, y: jnp.linalg.norm(x - y) ** 2
jax.grad(l2_dist_sq)(x, x) #[nan nan nan]

# This should be the same as l2_dist_sq
oracle_l2_dist_sq = lambda x, y: (x - y).dot(x - y)
jax.grad(oracle_l2_dist_sq)(x, x) #[0 0 0]"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1.</span>, <span class="pl-c1">2.</span>, <span class="pl-c1">3.</span>])
<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-s1">norm</span>)(<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">x</span>) <span class="pl-c">#[nan nan nan]</span>

<span class="pl-s1">l2_dist_sq</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">x</span>, <span class="pl-s1">y</span>: <span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">y</span>) <span class="pl-c1">**</span> <span class="pl-c1">2</span>
<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">l2_dist_sq</span>)(<span class="pl-s1">x</span>, <span class="pl-s1">x</span>) <span class="pl-c">#[nan nan nan]</span>

<span class="pl-c"># This should be the same as l2_dist_sq</span>
<span class="pl-s1">oracle_l2_dist_sq</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">x</span>, <span class="pl-s1">y</span>: (<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">y</span>).<span class="pl-en">dot</span>(<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">y</span>)
<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">oracle_l2_dist_sq</span>)(<span class="pl-s1">x</span>, <span class="pl-s1">x</span>) <span class="pl-c">#[0 0 0]</span></pre></div>
<p dir="auto">Of course, I highly doubt <code class="notranslate">jnp.linalg.norm</code> is the only instance of this "bug", but I suspect this is a very commonly used function. It took me fairly long to deduce that my <code class="notranslate">nan</code>s were coming from this gradient. I suppose a simple fix in this case would be something like</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="def norm(x, eps=1e-9):
    norm_sq = x.dot(x)
    return jnp.where(norm_sq &lt; eps, 0, jnp.sqrt(norm_sq))"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">norm</span>(<span class="pl-s1">x</span>, <span class="pl-s1">eps</span><span class="pl-c1">=</span><span class="pl-c1">1e-9</span>):
    <span class="pl-s1">norm_sq</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>.<span class="pl-en">dot</span>(<span class="pl-s1">x</span>)
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">where</span>(<span class="pl-s1">norm_sq</span> <span class="pl-c1">&lt;</span> <span class="pl-s1">eps</span>, <span class="pl-c1">0</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">norm_sq</span>))</pre></div>
<p dir="auto">Having said that, perhaps there's a more conventional way of achieving this?</p>