Open stuurver opened 1 year ago
Can confirm. I get the exact same error
will take a look thx for reporting. 😄
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:
jax
ortpu_driver_20221011
value we used in the notebook to a newer one.I have no clue. @charlielito you mind having a look at this?
I got this to work last night by doing the following:
0.3.25
)jax.tools.colab_tpu.setup_tpu()
Still would prefer to wait for @charlielito to weigh in before I make any changes to colab
Thanks this seems to work indeed!
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
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.
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.
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
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'