I also encouter this problem since a day or 2. Crashes everytime I run it with this error message:

UnfilteredStackTrace: AttributeError: 'UnspecifiedValue' object has no attribute '_parsed_pspec'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

The above exception was the direct cause of the following exception:

AttributeError Traceback (most recent call last)
/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py in threefry_seed(seed)
824 # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit
825 # inputs. We should avoid this.
--> 826 k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
827 return lax.concatenate([k1, k2], 0)
828

AttributeError: 'UnspecifiedValue' object has no attribute '_parsed_pspec'