<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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">PyTorch will soon add bf16 support on GPUs. Where is JAX on that?</p>
<p dir="auto">Will JAX soon add support for bf16 precision on GPUs?</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>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<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

f = lambda x: jnp.square(x).mean()
jf = jax.jit(f)

x = jax.random.uniform(jax.random.PRNGKey(0), shape=(8, 4))

# Start monitoring memory with htop
while True:
    try:
        # let this run for a bit, notice memory usage doesn't increase
        y = jax.hessian(f)(x)
    except KeyboardInterrupt:
        pass

while True:
    # let this run, notice memory usage starts increasing immediately
    # also, the memory usage doesn't appear to stop increasing ever
    y = jax.hessian(jf)(x)"><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">f</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">square</span>(<span class="pl-s1">x</span>).<span class="pl-en">mean</span>()
<span class="pl-s1">jf</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">f</span>)

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</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">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">8</span>, <span class="pl-c1">4</span>))

<span class="pl-c"># Start monitoring memory with htop</span>
<span class="pl-k">while</span> <span class="pl-c1">True</span>:
    <span class="pl-k">try</span>:
        <span class="pl-c"># let this run for a bit, notice memory usage doesn't increase</span>
        <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">hessian</span>(<span class="pl-s1">f</span>)(<span class="pl-s1">x</span>)
    <span class="pl-k">except</span> <span class="pl-v">KeyboardInterrupt</span>:
        <span class="pl-k">pass</span>

<span class="pl-k">while</span> <span class="pl-c1">True</span>:
    <span class="pl-c"># let this run, notice memory usage starts increasing immediately</span>
    <span class="pl-c"># also, the memory usage doesn't appear to stop increasing ever</span>
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">hessian</span>(<span class="pl-s1">jf</span>)(<span class="pl-s1">x</span>)</pre></div>
<p dir="auto">Also, note that the equivalent code with <code class="notranslate">jax.grad</code> (or <code class="notranslate">jax.jacfwd</code> for multidimensional functions) replacing <code class="notranslate">jax.hessian</code> <em>does not</em> cause a memory leak.</p>
<p dir="auto">As a result of this leak, my python process ends up getting killed by my OS (due to OOM), so I don't ever see any python or jax traces.</p>