jaxtyping not working with threads

[ayaka14732]

Typecheck cannot be performed in another thread:

import jax.numpy as np
from jaxtyping import Array, Float as F, jaxtyped
import threading
from typeguard import typechecked as typechecker

@jaxtyped
@typechecker
def add(x: F[Array, 'a b'], y: F[Array, 'a b']) -> F[Array, 'a b']:
    return x + y

def run():
    a = np.array([[1., 2.]])
    b = np.array([[2., 3.]])
    c = add(a, b)
    print(c)

thread = threading.Thread(target=run)
thread.start()
thread.join()

Error:

Exception in thread Thread-1 (run):
Traceback (most recent call last):
  File "/usr/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/home/ayaka/Projects/test/1.py", line 14, in run
    c = add(a, b)
  File "/home/ayaka/Projects/test/venv/lib/python3.10/site-packages/jaxtyping/decorator.py", line 31, in wrapper
    storage.memo_stack.append(({}, {}, {}))
AttributeError: '_thread._local' object has no attribute 'memo_stack'

Environment:

• Python 3.10.6
• jax 0.3.17
• jaxlib 0.3.15+cuda11.cudnn82
• jaxtyping 0.2.1

[patrick-kidger]

Should be fixed in #24!

[wookayin]

I'm still having a similar issue with jaxtyping==0.2.2 and jaxtyping==0.2.3 -- from a similar but different place. Reverting back to jaxtyping==0.2.0 has no issue.

Client code:

assert isinstance(observations, types.FloatArray)

Where it crashes:

File "...python3.10/site-packages/jaxtyping/array_types.py", line 153, in __instancecheck__
  153       if len(storage.memo_stack) == 0:   

[patrick-kidger]

Can you provide a full MWE that I can run?

[wookayin]

It's just making any isinstance call within a thread:

Here is a minimal working example that fails with jaxtyping==0.2.3:

import jax.numpy as jnp
from jaxtyping import Array, Float as F
import threading
from typeguard import typechecked as typechecker

FloatArray = F[jnp.ndarray, '...']

def run():
    a = jnp.array([[1., 2.]])
    assert isinstance(a, FloatArray)

thread = threading.Thread(target=run)
thread.start()
thread.join()

Error:

  File "test.py", line 10, in run
    assert isinstance(a, FloatArray)
  File ".../jaxtyping/array_types.py", line 153, in __instancecheck__
    if len(storage.memo_stack) == 0:
AttributeError: '_thread._local' object has no attribute 'memo_stack'

[patrick-kidger]

Thanks! This should be fixed in #27.