Repository: google/jax

Issue:
multi-thread jax.grad fails on jax.lax.scan
Calculating the gradient of a function in multiple independent threads sometimes causes jax to encounter errors of the form.

```
[...]
third_party/py/jax/core.py in pop(self, bottom)
    605   def pop(self, bottom: bool) -> None:
    606     if bottom:
--> 607       self.downward.pop()
    608     else:
    609       self.upward.pop()

IndexError: pop from empty list
```

In my original setup this happened only in 1% of the cases, here is a minimal example which reproduces the issue most of the time:

```
import concurrent
import jax
from jax import numpy as jnp

def unroll(param, sequence):
  def scan_f(prev_state, inputs):
    return prev_state, jax.nn.sigmoid(param * inputs)
  return jnp.sum(jax.lax.scan(scan_f, None, sequence)[1])

def run():
  return jax.grad(unroll)(jnp.array(1.0), jnp.array([1.0]))

# The more workers the more likely the issue appears. Using 1 worker works as expected.
n_workers = 20
with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as e:
  futures = []
  for _ in range(n_workers):
    futures.append(e.submit(run))
  unused_results = [f.result() for f in futures]
```


Discussion:
Thanks for finding this, and the excellent repro!

I suspect the `custom_jvp` on `jax.nn.sigmoid` may be an issue here; if I replace `jax.nn.sigmoid` with `jnp.sin` the issue goes away.

(By the way, I think `import concurrent` should be `import concurrent.futures`.)
I further suspect that [this state saving/restoration](https://github.com/google/jax/blob/f6221a663ef6c311e986aeffcef7403be09e8b15/jax/custom_derivatives.py#L61) is not safe.
Yes, that line is definitely not safe; it's swapping a module-level variable! We need to swap only thread-local variables.

I think I've got a fix.

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -635,6 +635,11 @@ def __init__(self) -> None:
     self.substack = [Sublevel(0)]
     self.initial_style = False
 
+  def set_state(self, other: 'TraceState') -> None:
+    self.trace_stack = other.trace_stack
+    self.sustack = other.substack[:]
+    self.initial_style = other.initial_style
+
   def copy(self):
     new = TraceState()
     new.trace_stack = self.trace_stack.copy()
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -58,11 +58,12 @@ def _memoize(thunk):
   saved_state = core.trace_state.copy()
   def memoized():
     if not cell:
-      prev_state, core.trace_state = core.trace_state, saved_state
+      prev_state = core.trace_state.copy()
+      core.trace_state.set_state(saved_state)
       try:
         cell.append(thunk())
       finally:
-        core.trace_state = prev_state
+        core.trace_state.set_state(prev_state)
     return cell[0]
   return memoized
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2517,6 +2517,24 @@ def f(x, y):
     expected = 12.
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_concurrent_initial_style(self):
+    # https://github.com/google/jax/issues/3843
+    def unroll(param, sequence):
+      def scan_f(prev_state, inputs):
+        return prev_state, jax.nn.sigmoid(param * inputs)
+      return jnp.sum(jax.lax.scan(scan_f, None, sequence)[1])
+
+    def run():
+      return jax.grad(unroll)(jnp.array(1.0), jnp.array([1.0]))
+
+    # we just don't want this to crash
+    n_workers = 20
+    with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as e:
+      futures = []
+      for _ in range(n_workers):
+        futures.append(e.submit(run))
+      _ = [f.result() for f in futures]
+
 
 class CustomVJPTest(jtu.JaxTestCase):