Fatal Python error: Aborted when running mnist.py example

[google-deepmind]

Fatal Python error: Aborted when running mnist.py example

[shakes76]

This is the solution for the case when using JAX with tensorflow datasets that use tensorflow for the preprocessing as XLA and TF both fight for the memory. Using the documentation highlighted in this issue above points out using tf.config.experimental.set_visible_devices([], "GPU")

[iobtl]

Following the configuration instructions in https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html worked.
Thank you! Will keep note of such a requirement in future. I don't have concurrent processes running so this may just be an internal system setup on my side.

[trevorcai]

2020-05-11 20:08:32.182307: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:236] failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED

This error is indicating that CUDA didn't initialize properly.
As you point out, one cause of this might be that many processes are concurrently trying to reserve GPU memory. Do you get this error when there are no other processes using your GPU? (Can be seen with nvidia-smi.)

The JAX equivalent of the above TF configuration can be found here:
https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

One other thing that might be worth trying is encasing all the TFDS logic in a with tf.device('cpu'): to make sure the TF executor & XLA aren't fighting in your setup.

I'm fairly confident that this is a configuration issue; the folks in google/jax may be more able to pinpoint the problem than me.