jaxtyping not working with threads about jaxtyping HOT 5 CLOSED

Thanks! This should be fixed in #27.

from jaxtyping.

Should be fixed in #24!

from jaxtyping.

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:   

from jaxtyping.

Can you provide a full MWE that I can run?

from jaxtyping.

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'

from jaxtyping.