You need to install flax 0.4.2. If you're using collab you just open a terminal (icon in the bottom left of the screen) and run:

    pip3 install flax==0.4.2