<p dir="auto">Calculating a Jacobian returns substantially different answers on CPU vs. TPU.</p>
<p dir="auto">In the example below, we calculate the Jacobian, sum its rows and compare it to the gradient of the sum (calculated using gradient). When running on CPU, the answers are indeed almost identical, however, when running on TPU we get a large discrepancy.</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

rng = jax.random.PRNGKey(0)
n = 1000
m = 10000

A = jax.random.normal(rng, (n, m))


def f(x):
  return jnp.dot(A, x)

g = jax.jit(jax.jacrev(f))
g_sum = jax.jit(jax.grad(lambda x: jnp.sum((f(x)))))

x0 = jax.random.normal(rng, (m,))

print(&quot;max numerical error = &quot;, jnp.max(jnp.abs(jnp.sum(g(x0), axis=0) - g_sum(x0))))"><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">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>)
<span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-c1">1000</span>
<span class="pl-s1">m</span> <span class="pl-c1">=</span> <span class="pl-c1">10000</span>

<span class="pl-v">A</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">rng</span>, (<span class="pl-s1">n</span>, <span class="pl-s1">m</span>))


<span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">x</span>):
  <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">dot</span>(<span class="pl-v">A</span>, <span class="pl-s1">x</span>)

<span class="pl-s1">g</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">jax</span>.<span class="pl-en">jacrev</span>(<span class="pl-s1">f</span>))
<span class="pl-s1">g_sum</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>((<span class="pl-en">f</span>(<span class="pl-s1">x</span>)))))

<span class="pl-s1">x0</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">rng</span>, (<span class="pl-s1">m</span>,))

<span class="pl-en">print</span>(<span class="pl-s">"max numerical error = "</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">max</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">abs</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-en">g</span>(<span class="pl-s1">x0</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>) <span class="pl-c1">-</span> <span class="pl-en">g_sum</span>(<span class="pl-s1">x0</span>))))</pre></div>
<p dir="auto">When running g on TPU we get</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; max numerical error =  0.21506119"><pre class="notranslate"><code class="notranslate">&gt;&gt;&gt; max numerical error =  0.21506119
</code></pre></div>
<p dir="auto">and on CPU we get</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; max numerical error =  0.0001449585"><pre class="notranslate"><code class="notranslate">&gt;&gt;&gt; max numerical error =  0.0001449585
</code></pre></div>
<p dir="auto">Follow-up from <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="559452832" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2161" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2161/hovercard" href="https://github.com/google/jax/issues/2161">#2161</a>:</p>
<p dir="auto">The default low precision is a bit of a footgun, at least when doing anything that isn't implementing a neural net layer. In my opinion, it would be much safer to use "highest" precision by default (which isn't that much slower) on float32 data. Neural net libraries, of course, can default to lower precision, so this really only effects users who directly use NumPy APIs or the @ infix operator.</p>