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

[stuurver]

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'

[cdj5492]

Can confirm. I get the exact same error

[nateraw]

will take a look thx for reporting. 😄

[nateraw]

Can confirm its not running as it was previously. I got it "running" by restarting runtime after all the installs, but it was very very slow. There was a new release of jax recently (Feb 16th), so I assume that's causing it. Maybe need to either:

• downgrade jax or
• perhaps just need to change the tpu_driver_20221011 value we used in the notebook to a newer one.

I have no clue. @charlielito you mind having a look at this?

[nateraw]

I got this to work last night by doing the following:

• Just using jax that comes with colab (version 0.3.25)
• changing setup to not specify build number, since they update it with the release jax.tools.colab_tpu.setup_tpu()

Still would prefer to wait for @charlielito to weigh in before I make any changes to colab

[stuurver]

Thanks this seems to work indeed!

[charlielito]

I got this to work last night by doing the following:

• Just using jax that comes with colab (version 0.3.25)
• changing setup to not specify build number, since they update it with the release jax.tools.colab_tpu.setup_tpu()

Still would prefer to wait for @charlielito to weigh in before I make any changes to colab

I was on vacation :p I'll take a look this week

[nateraw]

All good @charlielito I hope you enjoyed your time off ❤️ . No rush here - let me know if I should merge my hotfix or if you think there's a better way whenever you get the chance.

[tralala87]

For me it doesn't work.
I get:
RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.4.

[charlielito]

For me it doesn't work. I get: RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.4.

Yeah, the new release of jax/jaxlib/flax is not compatible with the TPU API we are using. I'll work on the versions needed to run it again