Flax doesn't work on google colab

It seems like flax just stopped working on google colab. Simply running
```
import jax
!pip install --quiet flax
import flax
```
yields the error
```
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
ipython 7.9.0 requires jedi>=0.10, which is not installed.
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
[](https://localhost:8080/#) in
1 import jax
2 get_ipython().system('pip install --quiet flax')
----> 3 import flax

4 frames
[/usr/local/lib/python3.8/dist-packages/flax/core/meta.py](https://localhost:8080/#) in Partitioned()
263 return self.replace(names=tuple(names))
264
--> 265 def get_partition_spec(self) -> jax.sharding.PartitionSpec:
266 """Returns the ``Partitionspec`` for this partitioned value."""
267 return jax.sharding.PartitionSpec(*self.names)

AttributeError: module 'jax.sharding' has no attribute 'PartitionSpec'
```
Can be solved by downgrading Flax to 0.6.4.