Add a warning if the user calls os.fork().

[hawkinsp]

Fixes #18852

[ScottTodd]

I'm seeing this warning even when not directly using JAX. Could this warning be more localized somehow?

Minimal repro: https://colab.research.google.com/gist/ScottTodd/476cfa0dd6620511ace15ce321b03a4e/colab-fork-jax-warning.ipynb

import os
import transformers
pid = os.fork()

RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  pid = os.fork()

Specific callstacks into os.fork() and this code from my project: https://gist.github.com/ScottTodd/981f1c8696887e0f9aef7fa57adbe84b

[hawkinsp]

Threading and os.fork() do not mix. You're playing with fire combining os.fork() and pytorch, in my opinion.

It's possible we could delay adding this check until you actually initialize jax, I suppose.

[jpienaar]

I think the main confusing part of this is that even when not using JAX (e.g., using PyTorch and it calls fork itself) one gets a warning about JAX. Delaying would at least avoid that happening if not using JAX at all, but seems would still trigger if one used both these libraries together.

[jakevdp]

jax is a pretty heavy import if you're not using it. Maybe the best fix is to ensure that transformers doesn't import jax unless it needs to?

[stellaraccident]

pytorch/pytorch#123954

[hawkinsp]

#20734 defers installing the warning until JAX initializes its internals. If you simply import but do not use jax, this should silence the warning.

[ScottTodd]

Thanks for the fast responses! This warning did help us uncover a few issues across various packages, at least :)