<p dir="auto">The <a href="https://gist.github.com/Dapid/b33919674f6a3cd9839473bfd2821e28">following example</a> explodes in memory usage when calling the jitted functions. The culprit seems to be the dynamic_slice. It happens calling with or without the GPU.</p>
<p dir="auto">Incidentally, it is also rather slow (twice slower than the plain Python version), so I welcome comments if there is a jax-friendlier way to rewrite it, but perhaps that is what jit should do?</p>
<p dir="auto">My system:</p>
<p dir="auto">Fedora 29, 31 GB RAM<br>
jax 0.1.36<br>
jaxlib 0.1.16, installed from the googleapis wheel.<br>
CUDA 10.0<br>
Nvidia driver 430.14.</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="from jax import tree_util
leaves, td = tree_util.tree_flatten( (1, (2, 3), [4], {'a': 42}, 99) )
td.num_leaves, td.num_nodes
# (6, 10)
td2 = tree_util.treedef_tuple(td.children())
td2.num_leaves, td2.num_nodes
# (0, 10)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">tree_util</span>
<span class="pl-s1">leaves</span>, <span class="pl-s1">td</span> <span class="pl-c1">=</span> <span class="pl-s1">tree_util</span>.<span class="pl-en">tree_flatten</span>( (<span class="pl-c1">1</span>, (<span class="pl-c1">2</span>, <span class="pl-c1">3</span>), [<span class="pl-c1">4</span>], {<span class="pl-s">'a'</span>: <span class="pl-c1">42</span>}, <span class="pl-c1">99</span>) )
<span class="pl-s1">td</span>.<span class="pl-s1">num_leaves</span>, <span class="pl-s1">td</span>.<span class="pl-s1">num_nodes</span>
<span class="pl-c"># (6, 10)</span>
<span class="pl-s1">td2</span> <span class="pl-c1">=</span> <span class="pl-s1">tree_util</span>.<span class="pl-en">treedef_tuple</span>(<span class="pl-s1">td</span>.<span class="pl-en">children</span>())
<span class="pl-s1">td2</span>.<span class="pl-s1">num_leaves</span>, <span class="pl-s1">td2</span>.<span class="pl-s1">num_nodes</span>
<span class="pl-c"># (0, 10)</span></pre></div>
<p dir="auto">Notice <code class="notranslate">td.num_leaves</code> is <code class="notranslate">6</code>, but <code class="notranslate">td2.num_leaves</code> is <code class="notranslate">0</code>.</p>