Merge pull request #18989 from hawkinsp:fork

PiperOrigin-RevId: 591286756
jax-ml · Dec 15, 2023 · eb595af · eb595af
2 parents c8b3567 + ec89e5e
commit eb595af
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
@@ -106,6 +106,16 @@
 )
 
 
+# Warn the user if they call fork(), because it's not going to go well for them.
+def _at_fork():
+  warnings.warn(
+    "os.fork() was called. os.fork() is incompatible with multithreaded code, "
+    "and JAX is multithreaded, so this will likely lead to a deadlock.",
+    RuntimeWarning, stacklevel=2)
+
+os.register_at_fork(before=_at_fork)
+
+
 # Backends