Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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.

The code to reproduce the problem can be found in deepmind/acme#233, and the original issue is deepmind/acme#235, another related issue may be #10763.

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.

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.

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.

I am using jax[cuda]==0.3.6.

If applicable, include full error messages/tracebacks.

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.

#0 0x00000000005b9fa3 in PyObject_Malloc ()
#1 0x00000000005a9c18 in PyType_GenericAlloc ()
#2 0x00007fff7165a5bd in xla::PyBuffer::Make(std::shared_ptr<xla::PyClient>, std::shared_ptr<xla::PjRtBuffer>, std::shared_ptr<xla::Traceback>) ()
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<pybind11::cpp_function::initialize<tensorflow::StatusOr<pybind11::object>, 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>(tensorflow::StatusOr<pybind11::object> (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}, tensorflow::StatusOr<pybind11::object>, 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>(pybind11::cpp_function::initialize<tensorflow::StatusOr<pybind11::object>, 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>(tensorflow::StatusOr<pybind11::object> (xla::PyClient::*)(pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&)::{lambda(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics)#1}&&, tensorflow::StatusOr<pybind11::object> (*)(xla::PyClient*, pybind11::handle, xla::PjRtDevice*, bool, xla::PjRtClient::HostBufferSemantics), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&)::{lambda(pybind11::detail::function_call&)#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 <RET> 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<pybind11::iterable>(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<pybind11::cpp_function::initialize<pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling>(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}, pybind11::object, xla::PyTreeDef const*, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling>(pybind11::cpp_function::initialize<pybind11::object, xla::PyTreeDef, pybind11::iterable, pybind11::name, pybind11::is_method, pybind11::sibling>(pybind11::object (xla::PyTreeDef::*)(pybind11::iterable) const, pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::{lambda(xla::PyTreeDef const*, pybind11::iterable)#1}&&, pybind11::object (*)(xla::PyTreeDef const*, pybind11::iterable), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&)::{lambda(pybind11::detail::function_call&)#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 <RET> 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=<optimised out>) at pthread_create.c:477
#68 0x00007ffff7eda163 in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:95

Is there any guess for why this happens in my case?

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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("cifar10", split="train", 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)))

gives me an error:

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->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.

But if I force TF to run on CPU with

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("cifar10", split="train", 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)))

Then it works!

Why does TF having access to the GPU affect JAX's ability to locate cuDNN?

Here's my shell.nix for complete reproducibility: https://gist.github.com/samuela/319059b88a46a994b4c10dfa718f379e

And here's a relevant comment on another issue: NixOS/nixpkgs#158186 (comment)

If applicable, include full error messages/tracebacks.