cannot import name 'optim' from 'flax' error

[hanshounsu]

Hello, I am currently trying the demo, synthesize_midi.ipynb, and when executing the below line

from music_spectrogram_diffusion.models.diffusion import models

the error from the title occurs:
ImportError: cannot import name 'optim' from 'flax' (/usr/local/lib/python3.8/dist-packages/flax/init.py)

I think the issue is due to update of t5x, and I'm struggling to solve the problem by myself. It would be grateful for some help..!

[cghawthorne]

Yeah, I think we no longer need to pin T5X to an old version now that Colab supports Python 3.8, but I'm getting some weird TPU errors when I try to use the newer packages. I'll let you know if I get it figured out, but feel free to experiment with unpinning some of the packages in the meantime.

[cghawthorne]

Reopening since this still isn't working (and I still haven't figured out a fix).