<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"> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">I am sorry for not having a dead-simple example that can be used to reproduce this issue at this point. It seems to happen in a specific setting where I tried to train a pixel-based agent with DeepMind's Acme library.<br>
The code to reproduce the problem can be found in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1247101583" data-permission-text="Title is private" data-url="https://github.com/deepmind/acme/issues/233" data-hovercard-type="issue" data-hovercard-url="/deepmind/acme/issues/233/hovercard" href="https://github.com/deepmind/acme/issues/233">deepmind/acme#233</a>, and the original issue is <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1248582718" data-permission-text="Title is private" data-url="https://github.com/deepmind/acme/issues/235" data-hovercard-type="issue" data-hovercard-url="/deepmind/acme/issues/235/hovercard" href="https://github.com/deepmind/acme/issues/235">deepmind/acme#235</a>, another related issue may be <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="4084809" data-permission-text="Title is private" data-url="https://github.com/google/jax/discussions/10763" data-hovercard-type="discussion" data-hovercard-url="/google/jax/discussions/10763/hovercard" href="https://github.com/google/jax/discussions/10763">#10763</a>.</p>
<p dir="auto">I will try to narrow done to the simplest example that can cause this, but at the moment I haven't successfully been able to do so.</p>
<p dir="auto">For this specific application, I am running JAX on both CPU and GPU concurrently on a single process with multithreading. In particular, I am working with a standard RL actor learner setup where the acting happens on the CPU and learning happens on the GPU.<br>
The actor loads the latest weights from the learner asynchronously with a ThreadPoolExecutor. I have 2 3080 on my machine, but I only use a single GPU for training.</p>
<p dir="auto">I am using <code class="notranslate">jax[cuda]==0.3.6</code>.</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">I set a gdb session to catch the segmentation fault, and here's the backtrace from the segfault thread. The segmentation fault seems to come from the xla_extension.so in jaxlib.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="#0  0x00000000005b9fa3 in PyObject_Malloc ()
#1  0x00000000005a9c18 in PyType_GenericAlloc ()
#2  0x00007fff7165a5bd in xla::PyBuffer::Make(std::shared_ptr&lt;xla::PyClient&gt;, std::shared_ptr&lt;xla::PjRtBuffer&gt;, std::shared_ptr&lt;xla::Traceback&gt;) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#3  0x00007fff71666228 in xla::PyClient::BufferFromPyval(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#4  0x00007fff713f608d in pybind11::cpp_function::initialize&lt;pybind11::cpp_function::initialize&lt;tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(tensorflow::StatusOr&lt;pybind11::object&gt; (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}, tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(pybind11::cpp_function::initialize&lt;tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(tensorflow::StatusOr&lt;pybind11::object&gt; (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}&amp;&amp;, tensorflow::StatusOr&lt;pybind11::object&gt; (*)(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(pybind11::detail::function_call&amp;)#3}::operator()(pybind11::detail::function_call) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#5  0x00007fff713edc7b in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#6  0x00000000005f3989 in PyCFunction_Call ()
#7  0x00000000005f3e1e in _PyObject_MakeTpCall ()
#8  0x000000000050b183 in ?? ()
#9  0x0000000000570035 in _PyEval_EvalFrameDefault ()
#10 0x00000000005f6836 in _PyFunction_Vectorcall ()
#11 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#12 0x000000000056939a in _PyEval_EvalCodeWithName ()
#13 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#14 0x00000000005f3547 in PyObject_Call ()
#15 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
--Type &lt;RET&gt; for more, q to quit, c to continue without paging--
#16 0x00000000005f6836 in _PyFunction_Vectorcall ()
#17 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#18 0x00000000005f6836 in _PyFunction_Vectorcall ()
#19 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#20 0x000000000056939a in _PyEval_EvalCodeWithName ()
#21 0x000000000050aaa0 in ?? ()
#22 0x000000000056c28c in _PyEval_EvalFrameDefault ()
#23 0x000000000056939a in _PyEval_EvalCodeWithName ()
#24 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#25 0x00000000005f3547 in PyObject_Call ()
#26 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
#27 0x00000000005006d4 in ?? ()
#28 0x0000000000510b02 in PyIter_Next ()
#29 0x00007fff7146c6a6 in pybind11::iterator::advance() () from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#30 0x00007fff7160c4d5 in pybind11::object xla::PyTreeDef::UnflattenImpl&lt;pybind11::iterable&gt;(pybind11::iterable) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#31 0x00007fff7160c98d in xla::PyTreeDef::Unflatten(pybind11::iterable) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#32 0x00007fff7160854b in pybind11::cpp_function::initialize&lt;pybind11::cpp_function::initialize&lt;pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}, pybind11::object, xla::PyTreeDef const*, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::cpp_function::initialize&lt;pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}&amp;&amp;, pybind11::object (*)(xla::PyTreeDef const*, pybind11::iterable), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(pybind11::detail::function_call&amp;)#3}::_FUN(pybind11::detail::function_call) () from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#33 0x00007fff713edc7b in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#34 0x00000000005f3989 in PyCFunction_Call ()
#35 0x00000000005f3e1e in _PyObject_MakeTpCall ()
#36 0x000000000050b183 in ?? ()
#37 0x0000000000570035 in _PyEval_EvalFrameDefault ()
#38 0x000000000056939a in _PyEval_EvalCodeWithName ()
--Type &lt;RET&gt; for more, q to quit, c to continue without paging--
#39 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#40 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#41 0x000000000056939a in _PyEval_EvalCodeWithName ()
#42 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#43 0x0000000000570035 in _PyEval_EvalFrameDefault ()
#44 0x00000000005f6836 in _PyFunction_Vectorcall ()
#45 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#46 0x000000000056939a in _PyEval_EvalCodeWithName ()
#47 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#48 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#49 0x00000000005f6836 in _PyFunction_Vectorcall ()
#50 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#51 0x00000000005f6836 in _PyFunction_Vectorcall ()
#52 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#53 0x00000000005f6836 in _PyFunction_Vectorcall ()
#54 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#55 0x00000000005f6836 in _PyFunction_Vectorcall ()
#56 0x00000000005f3547 in PyObject_Call ()
#57 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
#58 0x00000000005f6836 in _PyFunction_Vectorcall ()
#59 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#60 0x00000000005f6836 in _PyFunction_Vectorcall ()
#61 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#62 0x00000000005f6836 in _PyFunction_Vectorcall ()
#63 0x000000000050aa2c in ?? ()
#64 0x00000000005f3547 in PyObject_Call ()
#65 0x0000000000655a9c in ?? ()
#66 0x0000000000675738 in ?? ()
#67 0x00007ffff7da0609 in start_thread (arg=&lt;optimised out&gt;) at pthread_create.c:477
#68 0x00007ffff7eda163 in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:95"><pre class="notranslate"><code class="notranslate">#0  0x00000000005b9fa3 in PyObject_Malloc ()
#1  0x00000000005a9c18 in PyType_GenericAlloc ()
#2  0x00007fff7165a5bd in xla::PyBuffer::Make(std::shared_ptr&lt;xla::PyClient&gt;, std::shared_ptr&lt;xla::PjRtBuffer&gt;, std::shared_ptr&lt;xla::Traceback&gt;) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#3  0x00007fff71666228 in xla::PyClient::BufferFromPyval(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#4  0x00007fff713f608d in pybind11::cpp_function::initialize&lt;pybind11::cpp_function::initialize&lt;tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(tensorflow::StatusOr&lt;pybind11::object&gt; (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}, tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(pybind11::cpp_function::initialize&lt;tensorflow::StatusOr&lt;pybind11::object&gt;, xla::PyClient, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v&gt;(tensorflow::StatusOr&lt;pybind11::object&gt; (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}&amp;&amp;, tensorflow::StatusOr&lt;pybind11::object&gt; (*)(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;, pybind11::arg const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;, pybind11::arg_v const&amp;)::{lambda(pybind11::detail::function_call&amp;)#3}::operator()(pybind11::detail::function_call) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#5  0x00007fff713edc7b in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#6  0x00000000005f3989 in PyCFunction_Call ()
#7  0x00000000005f3e1e in _PyObject_MakeTpCall ()
#8  0x000000000050b183 in ?? ()
#9  0x0000000000570035 in _PyEval_EvalFrameDefault ()
#10 0x00000000005f6836 in _PyFunction_Vectorcall ()
#11 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#12 0x000000000056939a in _PyEval_EvalCodeWithName ()
#13 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#14 0x00000000005f3547 in PyObject_Call ()
#15 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
--Type &lt;RET&gt; for more, q to quit, c to continue without paging--
#16 0x00000000005f6836 in _PyFunction_Vectorcall ()
#17 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#18 0x00000000005f6836 in _PyFunction_Vectorcall ()
#19 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#20 0x000000000056939a in _PyEval_EvalCodeWithName ()
#21 0x000000000050aaa0 in ?? ()
#22 0x000000000056c28c in _PyEval_EvalFrameDefault ()
#23 0x000000000056939a in _PyEval_EvalCodeWithName ()
#24 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#25 0x00000000005f3547 in PyObject_Call ()
#26 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
#27 0x00000000005006d4 in ?? ()
#28 0x0000000000510b02 in PyIter_Next ()
#29 0x00007fff7146c6a6 in pybind11::iterator::advance() () from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#30 0x00007fff7160c4d5 in pybind11::object xla::PyTreeDef::UnflattenImpl&lt;pybind11::iterable&gt;(pybind11::iterable) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#31 0x00007fff7160c98d in xla::PyTreeDef::Unflatten(pybind11::iterable) const ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#32 0x00007fff7160854b in pybind11::cpp_function::initialize&lt;pybind11::cpp_function::initialize&lt;pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}, pybind11::object, xla::PyTreeDef const*, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::cpp_function::initialize&lt;pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling&gt;(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}&amp;&amp;, pybind11::object (*)(xla::PyTreeDef const*, pybind11::iterable), pybind11::name const&amp;, pybind11::is_method const&amp;, pybind11::sibling const&amp;)::{lambda(pybind11::detail::function_call&amp;)#3}::_FUN(pybind11::detail::function_call) () from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#33 0x00007fff713edc7b in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) ()
   from /home/yicheng/virtualenvs/orlb/lib/python3.8/site-packages/jaxlib/xla_extension.so
#34 0x00000000005f3989 in PyCFunction_Call ()
#35 0x00000000005f3e1e in _PyObject_MakeTpCall ()
#36 0x000000000050b183 in ?? ()
#37 0x0000000000570035 in _PyEval_EvalFrameDefault ()
#38 0x000000000056939a in _PyEval_EvalCodeWithName ()
--Type &lt;RET&gt; for more, q to quit, c to continue without paging--
#39 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#40 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#41 0x000000000056939a in _PyEval_EvalCodeWithName ()
#42 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#43 0x0000000000570035 in _PyEval_EvalFrameDefault ()
#44 0x00000000005f6836 in _PyFunction_Vectorcall ()
#45 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#46 0x000000000056939a in _PyEval_EvalCodeWithName ()
#47 0x00000000005f6a13 in _PyFunction_Vectorcall ()
#48 0x000000000056b0ae in _PyEval_EvalFrameDefault ()
#49 0x00000000005f6836 in _PyFunction_Vectorcall ()
#50 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#51 0x00000000005f6836 in _PyFunction_Vectorcall ()
#52 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#53 0x00000000005f6836 in _PyFunction_Vectorcall ()
#54 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#55 0x00000000005f6836 in _PyFunction_Vectorcall ()
#56 0x00000000005f3547 in PyObject_Call ()
#57 0x000000000056c8cd in _PyEval_EvalFrameDefault ()
#58 0x00000000005f6836 in _PyFunction_Vectorcall ()
#59 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#60 0x00000000005f6836 in _PyFunction_Vectorcall ()
#61 0x000000000056b1da in _PyEval_EvalFrameDefault ()
#62 0x00000000005f6836 in _PyFunction_Vectorcall ()
#63 0x000000000050aa2c in ?? ()
#64 0x00000000005f3547 in PyObject_Call ()
#65 0x0000000000655a9c in ?? ()
#66 0x0000000000675738 in ?? ()
#67 0x00007ffff7da0609 in start_thread (arg=&lt;optimised out&gt;) at pthread_create.c:477
#68 0x00007ffff7eda163 in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:95
</code></pre></div>
<p dir="auto">Is there any guess for why this happens in my case?</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.numpy as jnp
import tensorflow_datasets as tfds
from flax import linen as nn
from jax import random

# See https://github.com/tensorflow/tensorflow/issues/53831.
train_ds = tfds.load(&quot;cifar10&quot;, split=&quot;train&quot;, as_supervised=True)

model = nn.Conv(features=1, kernel_size=(3, 3), strides=(1, 1))
params = model.init(random.PRNGKey(123), jnp.zeros((1, 32, 32, 3)))"><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-k">import</span> <span class="pl-s1">tensorflow_datasets</span> <span class="pl-k">as</span> <span class="pl-s1">tfds</span>
<span class="pl-k">from</span> <span class="pl-s1">flax</span> <span class="pl-k">import</span> <span class="pl-s1">linen</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</span>

<span class="pl-c"># See https://github.com/tensorflow/tensorflow/issues/53831.</span>
<span class="pl-s1">train_ds</span> <span class="pl-c1">=</span> <span class="pl-s1">tfds</span>.<span class="pl-en">load</span>(<span class="pl-s">"cifar10"</span>, <span class="pl-s1">split</span><span class="pl-c1">=</span><span class="pl-s">"train"</span>, <span class="pl-s1">as_supervised</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Conv</span>(<span class="pl-s1">features</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">3</span>, <span class="pl-c1">3</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>, <span class="pl-c1">1</span>))
<span class="pl-s1">params</span> <span class="pl-c1">=</span> <span class="pl-s1">model</span>.<span class="pl-en">init</span>(<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">123</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">1</span>, <span class="pl-c1">32</span>, <span class="pl-c1">32</span>, <span class="pl-c1">3</span>)))</pre></div>
<p dir="auto">gives me an error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv = (f32[1,32,32,1]{2,1,3,0}, u8[0]{0}) custom-call(f32[1,32,32,3]{2,1,3,0} %copy.3, f32[3,3,3,1]{1,0,2,3} %copy.4), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io-&gt;b01f, custom_call_target=&quot;__cudnn$convForward&quot;, metadata={op_type=&quot;conv_general_dilated&quot; op_name=&quot;jit(conv_general_dilated)/conv_general_dilated[\n  batch_group_count=1\n  dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n  feature_group_count=1\n  lhs_dilation=(1, 1)\n  lhs_shape=(1, 32, 32, 3)\n  padding=((1, 1), (1, 1))\n  precision=None\n  preferred_element_type=None\n  rhs_dilation=(1, 1)\n  rhs_shape=(3, 3, 3, 1)\n  window_strides=(1, 1)\n]&quot; source_file=&quot;/nix/store/ys9bmmwpdqf3vlgxjvfy770qdk4dcf1n-python3.9-flax-0.3.6/lib/python3.9/site-packages/flax/linen/linear.py&quot; source_line=282}, backend_config=&quot;{\&quot;conv_result_scale\&quot;:1,\&quot;activation_mode\&quot;:\&quot;0\&quot;,\&quot;side_input_scale\&quot;:0}&quot;

Original error: UNIMPLEMENTED: DNN library is not found."><pre class="notranslate"><code class="notranslate">RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv = (f32[1,32,32,1]{2,1,3,0}, u8[0]{0}) custom-call(f32[1,32,32,3]{2,1,3,0} %copy.3, f32[3,3,3,1]{1,0,2,3} %copy.4), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io-&gt;b01f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="jit(conv_general_dilated)/conv_general_dilated[\n  batch_group_count=1\n  dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n  feature_group_count=1\n  lhs_dilation=(1, 1)\n  lhs_shape=(1, 32, 32, 3)\n  padding=((1, 1), (1, 1))\n  precision=None\n  preferred_element_type=None\n  rhs_dilation=(1, 1)\n  rhs_shape=(3, 3, 3, 1)\n  window_strides=(1, 1)\n]" source_file="/nix/store/ys9bmmwpdqf3vlgxjvfy770qdk4dcf1n-python3.9-flax-0.3.6/lib/python3.9/site-packages/flax/linen/linear.py" source_line=282}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: UNIMPLEMENTED: DNN library is not found.
</code></pre></div>
<p dir="auto">But if I force TF to run on CPU with</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf

tf.config.set_visible_devices([], 'GPU')

import jax.numpy as jnp
import tensorflow_datasets as tfds
from flax import linen as nn
from jax import random

# See https://github.com/tensorflow/tensorflow/issues/53831.
train_ds = tfds.load(&quot;cifar10&quot;, split=&quot;train&quot;, as_supervised=True)

model = nn.Conv(features=1, kernel_size=(3, 3), strides=(1, 1))
params = model.init(random.PRNGKey(123), jnp.zeros((1, 32, 32, 3)))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>

<span class="pl-s1">tf</span>.<span class="pl-s1">config</span>.<span class="pl-en">set_visible_devices</span>([], <span class="pl-s">'GPU'</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-k">import</span> <span class="pl-s1">tensorflow_datasets</span> <span class="pl-k">as</span> <span class="pl-s1">tfds</span>
<span class="pl-k">from</span> <span class="pl-s1">flax</span> <span class="pl-k">import</span> <span class="pl-s1">linen</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</span>

<span class="pl-c"># See https://github.com/tensorflow/tensorflow/issues/53831.</span>
<span class="pl-s1">train_ds</span> <span class="pl-c1">=</span> <span class="pl-s1">tfds</span>.<span class="pl-en">load</span>(<span class="pl-s">"cifar10"</span>, <span class="pl-s1">split</span><span class="pl-c1">=</span><span class="pl-s">"train"</span>, <span class="pl-s1">as_supervised</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Conv</span>(<span class="pl-s1">features</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">3</span>, <span class="pl-c1">3</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>, <span class="pl-c1">1</span>))
<span class="pl-s1">params</span> <span class="pl-c1">=</span> <span class="pl-s1">model</span>.<span class="pl-en">init</span>(<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">123</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">1</span>, <span class="pl-c1">32</span>, <span class="pl-c1">32</span>, <span class="pl-c1">3</span>)))</pre></div>
<p dir="auto">Then it works!</p>
<p dir="auto">Why does TF having access to the GPU affect JAX's ability to locate cuDNN?</p>
<p dir="auto">Here's my shell.nix for complete reproducibility: <a href="https://gist.github.com/samuela/319059b88a46a994b4c10dfa718f379e">https://gist.github.com/samuela/319059b88a46a994b4c10dfa718f379e</a><br>
And here's a relevant comment on another issue: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1124747477" data-permission-text="Title is private" data-url="https://github.com/NixOS/nixpkgs/issues/158186" data-hovercard-type="pull_request" data-hovercard-url="/NixOS/nixpkgs/pull/158186/hovercard?comment_id=1030486912&amp;comment_type=issue_comment" href="https://github.com/NixOS/nixpkgs/pull/158186#issuecomment-1030486912">NixOS/nixpkgs#158186 (comment)</a></p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>